{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ozc9I4X1C4nX"
      },
      "source": [
        "Copyright 2022 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "aIZExrKYCqRi"
      },
      "outputs": [],
      "source": [
        "#@title License\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4mdq2mi2Y7_N"
      },
      "source": [
        "# Sequential attention time tracking\n",
        "\n",
        "This directory hosts a Python 3 notebook that extends the work in [Yasuda et al. 2023](https://arxiv.org/abs/2209.14881). We run a comparative study using both low data regimes and low features regimes, with both Simple Sequential Attention and Intermediate(Ensembled) Sequential Attention to understand how we can both lower the data quantity as well as produce better and more consistent results."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JPlEvB8OJzyp"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ym5fXWFhMVxq"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "from random import shuffle\n",
        "import time\n",
        "from typing import List, Optional\n",
        "\n",
        "from matplotlib import pyplot as plt\n",
        "import matplotlib.patches as mpatches\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from sklearn.preprocessing import MinMaxScaler\n",
        "import tensorflow as tf"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vfZu-gluJ0xF"
      },
      "source": [
        "## Models\n",
        "\n",
        "\n",
        "For the purpose of this reproduction, multiple models will be defined as follows:\n",
        "  - `SequentialAttention`: the sequential attention block from the paper, with some minor changes to help with the experiments.\n",
        "  - `MLPModel`: simple neural network, with multiple fully connected layers.\n",
        "  - `SequentialAttentionMlp`: A fully connected network, with a Sequential Attention block before the dense layers. The block zeros out some of the features, based on the attention logits and the training progress.\n",
        "  - `SparseModel`: A fully connected network that takes into account the feature mask computed by Sequential Attention, and zeros out features before training or inference, as specified in the constructor.\n",
        "  - `LiaoLattyYangMask` and `LiaoLattyYangModel`: feature selection models, used for comparison with the paper of [Liao et al. 2021](https://arxiv.org/pdf/2010.13631.pdf). It is the same model used in the Sequential Attention paper for benchmarking.\n",
        "  - `SequentialLassoRegularizer` and `SequentialLassoModel`: Feature selection models based on Lasso, as defined by the Sequential Attention paper."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N1Tx-ZgZNsqd"
      },
      "outputs": [],
      "source": [
        "class SequentialAttention(tf.Module):\n",
        "  \"\"\"SequentialAttention module.\n",
        "  \n",
        "  Adapted from the Sequential Attention paper.\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      num_candidates,\n",
        "      num_candidates_to_select,\n",
        "      num_candidates_to_select_per_step=1,\n",
        "      start_percentage=0.05,\n",
        "      stop_percentage=1.0,\n",
        "      name='sequential_attention',\n",
        "      reset_weights=True,\n",
        "      preselected_candidate_indices: Optional[list[int]] = None,\n",
        "      **kwargs,\n",
        "  ):\n",
        "    \"\"\"Creates a new SequentialAttention module.\"\"\"\n",
        "\n",
        "    super(SequentialAttention, self).__init__(name=name, **kwargs)\n",
        "\n",
        "    assert num_candidates_to_select % num_candidates_to_select_per_step == 0, (\n",
        "        'num_candidates_to_select must be a multiple of '\n",
        "        'num_candidates_to_select_per_step.'\n",
        "    )\n",
        "\n",
        "    with self.name_scope:\n",
        "      self._num_candidates = num_candidates\n",
        "      self._num_candidates_to_select_per_step = (\n",
        "          num_candidates_to_select_per_step\n",
        "      )\n",
        "      self._num_steps = (\n",
        "          num_candidates_to_select // num_candidates_to_select_per_step\n",
        "      )\n",
        "      self._start_percentage = start_percentage\n",
        "      self._stop_percentage = stop_percentage\n",
        "      self._reset_weights = reset_weights\n",
        "\n",
        "      init_attention_weights = tf.random.normal(\n",
        "          shape=[num_candidates], stddev=0.00001, dtype=tf.float32\n",
        "      )\n",
        "      self._attention_weights = tf.Variable(\n",
        "          initial_value=lambda: init_attention_weights,\n",
        "          dtype=tf.float32,\n",
        "          name='attention_weights',\n",
        "      )\n",
        "\n",
        "      init_candidate_mask = [0.0] * num_candidates\n",
        "      self._preselected_candidate_indices_len = 0\n",
        "      if preselected_candidate_indices is not None:\n",
        "        self._preselected_candidate_indices_len = len(\n",
        "            preselected_candidate_indices\n",
        "        )\n",
        "        for index in preselected_candidate_indices:\n",
        "          assert index \u003e= 0\n",
        "          assert index \u003c num_candidates\n",
        "          init_candidate_mask[index] = 1.0\n",
        "\n",
        "      initial_weights = tf.constant(init_candidate_mask, dtype=tf.float32)\n",
        "      self.selected_features = tf.Variable(\n",
        "          initial_weights,\n",
        "          trainable=False,\n",
        "          name='selected_features',\n",
        "      )\n",
        "\n",
        "  @tf.Module.with_name_scope\n",
        "  def __call__(self, training_percentage):\n",
        "    \"\"\"Calculates attention weights for all candidates.\n",
        "\n",
        "    Args:\n",
        "      training_percentage: Percentage of training process that has been done.\n",
        "        This input argument should be between 0 and 1 and should be montonically\n",
        "        increasing.\n",
        "\n",
        "    Returns:\n",
        "      A vector of attention weights of size self._num_candidates. All the\n",
        "      weights\n",
        "      are between 0 and 1 and sum to 1.\n",
        "    \"\"\"\n",
        "    percentage = (training_percentage - self._start_percentage) / (\n",
        "        self._stop_percentage - self._start_percentage\n",
        "    )\n",
        "    curr_index = tf.cast(\n",
        "        tf.math.floor(percentage * self._num_steps), dtype=tf.float32\n",
        "    )\n",
        "    curr_index = tf.math.minimum(curr_index, self._num_steps - 1.0)\n",
        "\n",
        "    should_train = tf.less(curr_index, 0.0)\n",
        "\n",
        "    num_selected = tf.math.reduce_sum(self.selected_features)\n",
        "    should_select = tf.greater_equal(curr_index, num_selected)\n",
        "    _, new_indices = tf.math.top_k(\n",
        "        self._softmax_with_mask(\n",
        "            self._attention_weights, 1.0 - self.selected_features\n",
        "        ),\n",
        "        k=self._num_candidates_to_select_per_step,\n",
        "    )\n",
        "    new_indices = self._k_hot_mask(new_indices, self._num_candidates)\n",
        "    new_indices = tf.cond(\n",
        "        should_select,\n",
        "        lambda: new_indices,\n",
        "        lambda: tf.zeros(self._num_candidates),\n",
        "    )\n",
        "    select_op = self.selected_features.assign_add(new_indices)\n",
        "    init_attention_weights = tf.random.normal(\n",
        "        shape=[self._num_candidates], stddev=0.00001, dtype=tf.float32\n",
        "    )\n",
        "    should_reset = tf.logical_and(should_select, self._reset_weights)\n",
        "    new_weights = tf.cond(\n",
        "        should_reset,\n",
        "        lambda: init_attention_weights,\n",
        "        lambda: self._attention_weights,\n",
        "    )\n",
        "    reset_op = self._attention_weights.assign(new_weights)\n",
        "\n",
        "    with tf.control_dependencies([select_op, reset_op]):\n",
        "      candidates = 1.0 - self.selected_features\n",
        "      softmax = self._softmax_with_mask(self._attention_weights, candidates)\n",
        "      return tf.cond(\n",
        "          should_train,\n",
        "          lambda: tf.ones(self._num_candidates),\n",
        "          lambda: softmax + self.selected_features,\n",
        "      )\n",
        "\n",
        "  @tf.Module.with_name_scope\n",
        "  def _k_hot_mask(self, indices, depth, dtype=tf.float32):\n",
        "    return tf.math.reduce_sum(tf.one_hot(indices, depth, dtype=dtype), 0)\n",
        "\n",
        "  @tf.Module.with_name_scope\n",
        "  def _softmax_with_mask(self, logits, mask):\n",
        "    shifted_logits = logits - tf.math.reduce_max(logits)\n",
        "    exp_shifted_logits = tf.math.exp(shifted_logits)\n",
        "    masked_exp_shifted_logits = tf.multiply(exp_shifted_logits, mask)\n",
        "    return tf.math.divide_no_nan(\n",
        "        masked_exp_shifted_logits, tf.math.reduce_sum(masked_exp_shifted_logits)\n",
        "    )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2EVv1tz1N7iH"
      },
      "outputs": [],
      "source": [
        "class MLPModel(tf.keras.Model):\n",
        "  \"\"\"MLP model.\n",
        "  \n",
        "  Fully connected model, used to evaluate the performance of a masked dataset.\n",
        "  The input to this model is already masked, so it acts as a simple list of Dense layers.\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      layer_sequence=None,\n",
        "      is_classification=True,\n",
        "      num_classes=None,\n",
        "      learning_rate=0.001,\n",
        "      decay_steps=500,\n",
        "      decay_rate=0.8,\n",
        "      alpha=0,\n",
        "      batch_norm=True,\n",
        "  ):\n",
        "    \"\"\"Initialize the model.\"\"\"\n",
        "\n",
        "    super().__init__()\n",
        "\n",
        "    if batch_norm:\n",
        "      self.batch_norm_layer = tf.keras.layers.BatchNormalization()\n",
        "    self.batch_norm = batch_norm\n",
        "\n",
        "    mlp_sequence = [\n",
        "        tf.keras.layers.Dense(\n",
        "            dim, activation=tf.keras.layers.LeakyReLU(alpha=alpha)\n",
        "        )\n",
        "        for dim in layer_sequence\n",
        "    ]\n",
        "    self.mlp_model = tf.keras.Sequential(mlp_sequence)\n",
        "    if is_classification:\n",
        "      self.mlp_predictor = tf.keras.layers.Dense(\n",
        "          num_classes, activation=\"softmax\", dtype=\"float32\"\n",
        "      )\n",
        "    else:\n",
        "      self.mlp_predictor = tf.keras.layers.Dense(1, dtype=\"float32\")\n",
        "\n",
        "    lr = tf.keras.optimizers.schedules.ExponentialDecay(\n",
        "        learning_rate,\n",
        "        decay_steps=decay_steps,\n",
        "        decay_rate=decay_rate,\n",
        "        staircase=False,\n",
        "    )\n",
        "    self.optimizer = tf.keras.optimizers.Adam(learning_rate=lr)\n",
        "\n",
        "  def call(self, inputs, training=False):\n",
        "    if self.batch_norm:\n",
        "      inputs = self.batch_norm_layer(inputs, training=training)\n",
        "    inputs = tf.multiply(inputs, self.selected_features)\n",
        "    representation = self.mlp_model(inputs)  # other layers\n",
        "    prediction = self.mlp_predictor(representation)\n",
        "    return prediction"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PUni4kdcN2-P"
      },
      "outputs": [],
      "source": [
        "class SequentialAttentionMlp(MLPModel):\n",
        "  \"\"\"SequentialAttentionMlp.\n",
        "  \n",
        "  A fully connected neural network preceded by a Sequential Attention block.\n",
        "  During training, the Sequential Attention block learns the optimal mask.\n",
        "  During inference, the Sequential Attention masks the input according to the learned mask.\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      num_train_steps=20,\n",
        "      feature_mask=None,\n",
        "      num_candidates_to_select=50,\n",
        "      num_candidates=784,\n",
        "      preselected_candidate_indices: Optional[list[int]] = None,\n",
        "      **kwargs,\n",
        "  ):\n",
        "    super(SequentialAttentionMlp, self).__init__(**kwargs)\n",
        "\n",
        "    self.sa = SequentialAttention(\n",
        "        num_candidates=num_candidates,\n",
        "        num_candidates_to_select=num_candidates_to_select,\n",
        "        preselected_candidate_indices=preselected_candidate_indices,\n",
        "    )\n",
        "    self.num_train_steps = num_train_steps\n",
        "    self.iterations = tf.Variable(\n",
        "        0,\n",
        "        trainable=False,\n",
        "        dtype=tf.float32,\n",
        "    )\n",
        "    self.feature_mask = feature_mask\n",
        "\n",
        "  @tf.function\n",
        "  def call(self, inputs, training=False):\n",
        "    if self.feature_mask is not None:\n",
        "      inputs = tf.math.multiply(inputs, self.feature_mask)\n",
        "\n",
        "    self.iterations.assign_add(1)\n",
        "    if self.batch_norm:\n",
        "      inputs = self.batch_norm_layer(inputs, training=training)\n",
        "\n",
        "    training_percentage = (\n",
        "        self.iterations / self.num_train_steps\n",
        "    )  # e.g., 0.5 is 50% of the training process\n",
        "    attention_weights = self.sa(training_percentage)\n",
        "\n",
        "    inputs = tf.math.multiply(inputs, attention_weights)\n",
        "    representation = self.mlp_model(inputs)  # other layers\n",
        "    prediction = self.mlp_predictor(representation)\n",
        "\n",
        "    return prediction"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SyqHDt9wTP1B"
      },
      "outputs": [],
      "source": [
        "class SparseModel(MLPModel):\n",
        "  \"\"\"SparseModel.\n",
        "  \n",
        "  Acts as a fully connected neural network, with the difference that it masks (zeros out) some features before training or inference, according to a mask specified in the constructor.\n",
        "  In comparison to SequentialAttentionMlp, this model does not learn a feature mask. It only operates with a predefined mask, that is applied to the input.\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self, selected_features, **kwargs):\n",
        "    \"\"\"Initialize the model.\"\"\"\n",
        "\n",
        "    super(SparseModel, self).__init__(**kwargs)\n",
        "    self.selected_features = selected_features\n",
        "\n",
        "  def call(self, inputs, training=False):\n",
        "    if self.batch_norm:\n",
        "      inputs = self.batch_norm_layer(inputs, training=training)\n",
        "    inputs = tf.multiply(inputs, self.selected_features)\n",
        "    representation = self.mlp_model(inputs)  # other layers\n",
        "    prediction = self.mlp_predictor(representation)\n",
        "    return prediction"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V7GvKMT0cSEt"
      },
      "outputs": [],
      "source": [
        "class LiaoLattyYangMask(tf.Module):\n",
        "  \"\"\"A feature selection algorithm based on attention mechanism.\"\"\"\n",
        "\n",
        "  def __init__(self, num_inputs, name='liao_latty_yang_mask', **kwargs):\n",
        "    super(LiaoLattyYangMask, self).__init__(name=name, **kwargs)\n",
        "\n",
        "    mlp_sequence = [\n",
        "        tf.keras.layers.Dense(\n",
        "            dim, activation=tf.keras.layers.LeakyReLU(alpha=0.2)\n",
        "        )\n",
        "        for dim in [128, 64, num_inputs]\n",
        "    ]\n",
        "    self.mlp_model = tf.keras.Sequential(mlp_sequence)\n",
        "\n",
        "  def __call__(self, inputs):\n",
        "    nonlinear = self.mlp_model(inputs)\n",
        "    batch_size = tf.cast(tf.shape(inputs)[0], tf.float32)\n",
        "    logits = tf.reduce_sum(nonlinear, axis=0) / batch_size\n",
        "    return tf.nn.softmax(logits)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HS2bgmNCcOXH"
      },
      "outputs": [],
      "source": [
        "class LiaoLattyYangModel(MLPModel):\n",
        "  \"\"\"MLP with Liao-Latty-Yang.\"\"\"\n",
        "\n",
        "  def __init__(self, num_inputs, **kwargs):\n",
        "    \"\"\"Initialize the model.\"\"\"\n",
        "\n",
        "    super(LiaoLattyYangModel, self).__init__(**kwargs)\n",
        "\n",
        "    self.lly = LiaoLattyYangMask(num_inputs=num_inputs)\n",
        "\n",
        "  def call(self, inputs, training=False, return_attention=False):\n",
        "    if self.batch_norm:\n",
        "      inputs = self.batch_norm_layer(inputs, training=training)\n",
        "    attention_weights = self.lly(inputs)\n",
        "    if return_attention:\n",
        "      return attention_weights\n",
        "    inputs = tf.multiply(inputs, attention_weights)\n",
        "    representation = self.mlp_model(inputs)  # other layers\n",
        "    prediction = self.mlp_predictor(representation)\n",
        "    return prediction"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f3HJECwoezxs"
      },
      "outputs": [],
      "source": [
        "class SequentialLassoRegularizer(tf.Module):\n",
        "  \"\"\"Sequential Lasso for feature selection.\"\"\"\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      num_inputs,\n",
        "      num_inputs_to_select,\n",
        "      num_inputs_to_select_per_step=1,\n",
        "      preselected_features=None,\n",
        "      group_lasso_scale=0.01,\n",
        "      name='sequential_lasso_regularizer',\n",
        "      **kwargs,\n",
        "  ):\n",
        "    super(SequentialLassoRegularizer, self).__init__(name=name, **kwargs)\n",
        "\n",
        "    assert num_inputs_to_select % num_inputs_to_select_per_step == 0\n",
        "\n",
        "    self.group_lasso_scale = group_lasso_scale\n",
        "\n",
        "    with self.name_scope:\n",
        "      self._num_inputs = num_inputs\n",
        "      self._num_inputs_to_select = num_inputs_to_select\n",
        "      self._num_inputs_to_select_per_step = num_inputs_to_select_per_step\n",
        "      if preselected_features is not None:\n",
        "        num_preselected_features = len(preselected_features)\n",
        "        preselected_features = tf.math.reduce_sum(\n",
        "            tf.one_hot(preselected_features, num_inputs, dtype=tf.float32), 0\n",
        "        )\n",
        "        self.selected_features = tf.Variable(\n",
        "            preselected_features, trainable=False, name='selected_features'\n",
        "        )\n",
        "        self._num_inputs_to_select -= num_preselected_features\n",
        "      else:\n",
        "        self.selected_features = tf.Variable(\n",
        "            tf.zeros(num_inputs), trainable=False, name='selected_features'\n",
        "        )\n",
        "      self._num_selected_features = tf.Variable(\n",
        "          0, trainable=False, name='num_selected_features'\n",
        "      )\n",
        "\n",
        "      self.selected_features_history = tf.Variable(\n",
        "          initial_value=tf.zeros(\n",
        "              shape=self._num_inputs_to_select, dtype=tf.int32\n",
        "          ),\n",
        "          trainable=False,\n",
        "          dtype=tf.int32,\n",
        "          name='selected_features_history',\n",
        "      )\n",
        "\n",
        "  @tf.Module.with_name_scope\n",
        "  def __call__(self, x, training_percentage):\n",
        "    start_percentage = 0.1\n",
        "    end_percentage = 0.999\n",
        "    effective_percentage = (training_percentage - start_percentage) / (\n",
        "        end_percentage - start_percentage\n",
        "    )\n",
        "\n",
        "    curr_phase = tf.cast(\n",
        "        tf.math.ceil(\n",
        "            effective_percentage\n",
        "            * (\n",
        "                self._num_inputs_to_select\n",
        "                // self._num_inputs_to_select_per_step\n",
        "            )\n",
        "        ),\n",
        "        tf.int32,\n",
        "    )\n",
        "\n",
        "    new_phase = tf.greater(\n",
        "        curr_phase,\n",
        "        self._num_selected_features // self._num_inputs_to_select_per_step,\n",
        "        name='new_phase',\n",
        "    )\n",
        "    done_phase = tf.greater(\n",
        "        curr_phase,\n",
        "        self._num_inputs_to_select // self._num_inputs_to_select_per_step,\n",
        "        name='done_phase',\n",
        "    )\n",
        "    should_select_feature = tf.math.logical_and(\n",
        "        new_phase, tf.math.logical_not(done_phase)\n",
        "    )\n",
        "\n",
        "    unselected_features = tf.ones(self._num_inputs) - self.selected_features\n",
        "    _, indices = tf.math.top_k(\n",
        "        tf.multiply(x, unselected_features),\n",
        "        k=self._num_inputs_to_select_per_step,\n",
        "    )\n",
        "    new_features = tf.math.reduce_sum(tf.one_hot(indices, self._num_inputs), 0)\n",
        "    new_features = tf.cond(\n",
        "        should_select_feature,\n",
        "        lambda: new_features,\n",
        "        lambda: tf.zeros(self._num_inputs),\n",
        "    )\n",
        "    selected_features = self.selected_features.assign_add(new_features)\n",
        "    unselected_features = tf.ones(self._num_inputs) - selected_features\n",
        "\n",
        "    # update history\n",
        "    support = tf.one_hot(\n",
        "        tf.range(\n",
        "            self._num_selected_features,\n",
        "            self._num_selected_features + self._num_inputs_to_select_per_step,\n",
        "        ),\n",
        "        self._num_inputs_to_select,\n",
        "        dtype=tf.int32,\n",
        "    )\n",
        "    indices = tf.reshape(\n",
        "        indices, shape=[self._num_inputs_to_select_per_step, 1]\n",
        "    )\n",
        "    history = tf.cond(\n",
        "        should_select_feature,\n",
        "        lambda: tf.reduce_sum(support * indices, 0),\n",
        "        lambda: tf.zeros(self._num_inputs_to_select, dtype=tf.int32),\n",
        "    )\n",
        "    history = self.selected_features_history.assign_add(history)\n",
        "\n",
        "    # update num_selected_features\n",
        "    with tf.control_dependencies([history]):  # force update\n",
        "      num_new_features = tf.cond(\n",
        "          should_select_feature,\n",
        "          lambda: self._num_inputs_to_select_per_step,\n",
        "          lambda: 0,\n",
        "      )\n",
        "      num_selected_features = self._num_selected_features.assign_add(\n",
        "          num_new_features\n",
        "      )\n",
        "    with tf.control_dependencies([num_selected_features]):  # force update\n",
        "      return self.group_lasso_scale * tf.reduce_sum(\n",
        "          tf.multiply(x, unselected_features)\n",
        "      )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ymLxflVDetuf"
      },
      "outputs": [],
      "source": [
        "class SequentialLASSOModel(MLPModel):\n",
        "  \"\"\"MLP with Sequential Lasso.\"\"\"\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      num_inputs,\n",
        "      num_inputs_to_select,\n",
        "      num_train_steps,\n",
        "      num_inputs_to_select_per_step=1,\n",
        "      layer_sequence=[67],\n",
        "      alpha=0,\n",
        "      group_lasso_scale=0.01,\n",
        "      **kwargs,\n",
        "  ):\n",
        "    \"\"\"Initialize the model.\"\"\"\n",
        "\n",
        "    super(SequentialLASSOModel, self).__init__(\n",
        "        layer_sequence=layer_sequence, **kwargs\n",
        "    )\n",
        "\n",
        "    # first layer\n",
        "    init_kernel = tf.random.normal(\n",
        "        shape=[num_inputs, layer_sequence[0]], stddev=0.001, dtype=tf.float32\n",
        "    )\n",
        "    self.kernel0 = tf.Variable(\n",
        "        initial_value=lambda: init_kernel, dtype=tf.float32, name='kernel'\n",
        "    )\n",
        "    self.bias0 = tf.Variable(\n",
        "        initial_value=tf.zeros([layer_sequence[0]]), dtype=tf.float32\n",
        "    )\n",
        "    # other layers\n",
        "    mlp_sequence = []\n",
        "    self.mlp_model = tf.keras.Sequential(mlp_sequence)\n",
        "\n",
        "    self.seql = SequentialLassoRegularizer(\n",
        "        num_inputs=num_inputs,\n",
        "        num_inputs_to_select=num_inputs_to_select,\n",
        "        num_inputs_to_select_per_step=num_inputs_to_select_per_step,\n",
        "        group_lasso_scale=group_lasso_scale,\n",
        "    )\n",
        "    self.num_train_steps = num_train_steps\n",
        "\n",
        "  def call(self, inputs, training=False):\n",
        "    if self.batch_norm:\n",
        "      inputs = self.batch_norm_layer(inputs, training=training)\n",
        "    inputs = tf.linalg.matmul(inputs, self.kernel0) + self.bias0  # first layer\n",
        "    representation = self.mlp_model(inputs)  # other layers\n",
        "    prediction = self.mlp_predictor(representation)\n",
        "    return prediction\n",
        "\n",
        "  def train_step(self, inputs):\n",
        "    \"\"\"Custom train step.\"\"\"\n",
        "    training_percentage = self.optimizer.iterations / self.num_train_steps\n",
        "\n",
        "    with tf.GradientTape() as tape:\n",
        "      x, y = inputs\n",
        "      y_pred = self.call(x, training=True)\n",
        "      norms = tf.norm(self.kernel0, axis=1)\n",
        "      reg = self.seql(norms, training_percentage=training_percentage)\n",
        "      loss = self.compute_loss(x, y, y_pred) + reg\n",
        "    gradients = tape.gradient(loss, self.trainable_variables)\n",
        "    self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))\n",
        "    self.compiled_metrics.update_state(y, y_pred)\n",
        "    return {m.name: m.result() for m in self.metrics}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IEAbzJY_J2p-"
      },
      "source": [
        "## Dataset\n",
        "\n",
        "We ran experiments on 2 classification datasets, one image based (mnist), and a tabular one (Activity). We chose classification to provide easier to understand metrics (test set accuracies), and we picked both image and tabular to understand the feature to feature interactions, as well as the difference between images (where the position of a feature matters) vs tabular (where the ordering of features is irrelevant). Also, those 2 datasets were chosen to be consistent with the original paper (Sequential Attention for Feature Selection)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6WBAKWELMxwH"
      },
      "outputs": [],
      "source": [
        "def load_mnist(subsample_size: Optional[float] = 1.0):\n",
        "  \"\"\"Loads the Fashion dataset, adapted from: https://github.com/lasso-net/lassonet/blob/master/experiments/data_utils.py.\"\"\"\n",
        "\n",
        "  (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n",
        "\n",
        "  x_train = x_train.reshape((-1, x_train.shape[1] * x_train.shape[2])).astype(\n",
        "      np.float32\n",
        "  )\n",
        "  x_test = x_test.reshape((-1, x_test.shape[1] * x_test.shape[2])).astype(\n",
        "      np.float32\n",
        "  )\n",
        "\n",
        "  x = np.concatenate((x_train, x_test))\n",
        "  x = (x - x.min()) / (x.max() - x.min())\n",
        "  x_train = x[: len(y_train)]\n",
        "  x_test = x[len(y_train) :]\n",
        "\n",
        "  train, test = (x_train, y_train), (x_test, y_test)\n",
        "\n",
        "  x, y = train\n",
        "  x_y = np.concatenate([x, y.reshape((len(y), 1))], axis=1)\n",
        "  sampled = x_y[\n",
        "      np.random.choice(len(x_y), int(subsample_size * len(x_y)), replace=False),\n",
        "      :,\n",
        "  ]\n",
        "  train = (sampled[:, :-1], sampled[:, -1])\n",
        "\n",
        "  is_classification = True\n",
        "  num_classes = 10\n",
        "\n",
        "  x_train = pd.DataFrame(train[0])\n",
        "  x_test = pd.DataFrame(test[0])\n",
        "  y_train = pd.DataFrame(train[1], dtype=np.int32).iloc[:, 0]\n",
        "  y_test = pd.DataFrame(test[1], dtype=np.int32).iloc[:, 0]\n",
        "\n",
        "  print(\"Data loaded...\")\n",
        "  print(\"Data shapes:\")\n",
        "  print(x_train.shape, y_train.shape)\n",
        "  print(x_test.shape, y_test.shape)\n",
        "\n",
        "  return (x_train, x_test, y_train, y_test, is_classification, num_classes)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cHSPEvnT6MQc"
      },
      "outputs": [],
      "source": [
        "# For the Activity files upload.\n",
        "# Download first from https://www.kaggle.com/datasets/uciml/human-activity-recognition-with-smartphones.\n",
        "from google.colab import files\n",
        "\n",
        "uploaded = files.upload()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NLriiAM26P1f"
      },
      "outputs": [],
      "source": [
        "def load_activity(subsample_size: Optional[float] = 1.0):\n",
        "  \"\"\"Loads the Activity dataset, adapted from: https://github.com/lasso-net/lassonet/blob/master/experiments/data_utils.py.\"\"\"\n",
        "\n",
        "  cache_filepath_train_x = \"activity_X_train.txt\"\n",
        "  cache_filepath_train_y = \"activity_y_train.txt\"\n",
        "  cache_filepath_test_x = \"activity_X_test.txt\"\n",
        "  cache_filepath_test_y = \"activity_y_test.txt\"\n",
        "  with open(cache_filepath_train_x, \"r\") as fp:\n",
        "    x_train = np.genfromtxt(fp.readlines(), encoding=\"UTF-8\")\n",
        "  with open(cache_filepath_test_x, \"r\") as fp:\n",
        "    x_test = np.genfromtxt(fp.readlines(), encoding=\"UTF-8\")\n",
        "  with open(cache_filepath_train_y, \"r\") as fp:\n",
        "    y_train = np.genfromtxt(fp.readlines(), encoding=\"UTF-8\")\n",
        "  with open(cache_filepath_test_y, \"r\") as fp:\n",
        "    y_test = np.genfromtxt(fp.readlines(), encoding=\"UTF-8\")\n",
        "\n",
        "  x = MinMaxScaler(feature_range=(0, 1)).fit_transform(\n",
        "      np.concatenate((x_train, x_test))\n",
        "  )\n",
        "  x_train = x[: len(y_train)]\n",
        "  x_test = x[len(y_train) :]\n",
        "\n",
        "  x_y = np.concatenate([x_train, y_train.reshape((len(y_train), 1))], axis=1)\n",
        "  sampled = x_y[\n",
        "      np.random.choice(len(x_y), int(subsample_size * len(x_y)), replace=False),\n",
        "      :,\n",
        "  ]\n",
        "  (x_train, y_train) = (sampled[:, :-1], sampled[:, -1])\n",
        "\n",
        "  print(\"Data loaded...\")\n",
        "  print(\"Data shapes:\")\n",
        "  print(x_train.shape, y_train.shape)\n",
        "  print(x_test.shape, y_test.shape)\n",
        "\n",
        "  is_classification = True\n",
        "  num_classes = 6\n",
        "\n",
        "  x_train = pd.DataFrame(x_train)\n",
        "  x_test = pd.DataFrame(x_test)\n",
        "  y_train = pd.DataFrame(y_train - 1, dtype=np.int32).iloc[:, 0]\n",
        "  y_test = pd.DataFrame(y_test - 1, dtype=np.int32).iloc[:, 0]\n",
        "\n",
        "  return (x_train, x_test, y_train, y_test, is_classification, num_classes)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "89qRTSkjMzpB"
      },
      "outputs": [],
      "source": [
        "def get_dataset(\n",
        "    val_ratio,\n",
        "    batch_size,\n",
        "    dataset=\"mnist\",\n",
        "    subsample_size: Optional[float] = 1.0,\n",
        "):\n",
        "  \"\"\"Get datasets split into training, validation, and test datasets.\"\"\"\n",
        "  # Load data.\n",
        "  if dataset == \"mice\":\n",
        "    (x_train, x_test, y_train, y_test, is_classification, num_classes) = (\n",
        "        load_mice()\n",
        "    )\n",
        "  elif dataset == \"activity\":\n",
        "    (x_train, x_test, y_train, y_test, is_classification, num_classes) = (\n",
        "        load_activity(subsample_size=subsample_size)\n",
        "    )\n",
        "  else:\n",
        "    (x_train, x_test, y_train, y_test, is_classification, num_classes) = (\n",
        "        load_mnist(subsample_size=subsample_size)\n",
        "    )\n",
        "\n",
        "  # Tensorflow data transform functions\n",
        "  if is_classification:\n",
        "\n",
        "    def transform(x, y):\n",
        "      x = tf.cast(x, dtype=tf.float32)\n",
        "      return x, tf.one_hot(y, num_classes)\n",
        "\n",
        "  else:\n",
        "\n",
        "    def transform(x, y):\n",
        "      x = tf.cast(x, dtype=tf.float32)\n",
        "      y = tf.cast(y, dtype=tf.float32)\n",
        "      return x, y\n",
        "\n",
        "  # Shuffle training data\n",
        "  idx = np.random.permutation(x_train.index)\n",
        "  val_size = int(np.size(idx) * val_ratio)\n",
        "  x_val = x_train.reindex(idx[-val_size:])\n",
        "  y_val = y_train.reindex(idx[-val_size:])\n",
        "  x_train = x_train.reindex(idx[:-val_size])\n",
        "  y_train = y_train.reindex(idx[:-val_size])\n",
        "\n",
        "  # Construct tf dataset\n",
        "  assert batch_size \u003c= x_train.shape[0], (\n",
        "      f\"Batch size {batch_size} is larger than training data size\"\n",
        "      f\" {x_train.shape[0]}.\"\n",
        "  )\n",
        "  ds_train = tf.data.Dataset.from_tensor_slices(\n",
        "      (x_train.values, y_train.T.values)\n",
        "  )\n",
        "  ds_train = ds_train.map(transform).shuffle(100, reshuffle_each_iteration=True)\n",
        "  ds_train = ds_train.batch(batch_size, drop_remainder=True)\n",
        "\n",
        "  ds_val = tf.data.Dataset.from_tensor_slices((x_val.values, y_val.T.values))\n",
        "  ds_val = ds_val.map(transform)\n",
        "  ds_val = ds_val.batch(batch_size, drop_remainder=False)\n",
        "\n",
        "  ds_test = tf.data.Dataset.from_tensor_slices((x_test.values, y_test.T.values))\n",
        "  ds_test = ds_test.map(transform)\n",
        "  ds_test = ds_test.batch(batch_size, drop_remainder=False)\n",
        "\n",
        "  return {\n",
        "      \"x_train\": x_train,\n",
        "      \"y_train\": y_train,\n",
        "      \"x_val\": x_val,\n",
        "      \"y_val\": y_val,\n",
        "      \"x_test\": x_test,\n",
        "      \"y_test\": y_test,\n",
        "      \"ds_train\": ds_train,\n",
        "      \"ds_val\": ds_val,\n",
        "      \"ds_test\": ds_test,\n",
        "      \"num_features\": len(x_train.columns),\n",
        "      \"is_classification\": is_classification,\n",
        "      \"num_classes\": num_classes,\n",
        "  }"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0PPw-6uRxB5_"
      },
      "outputs": [],
      "source": [
        "def load_dataset(\n",
        "    num_epochs_select: int = 10,\n",
        "    batch_size: int = 256,\n",
        "    dataset=\"mnist\",\n",
        "    subsample_size: Optional[float] = 1.0,\n",
        "):\n",
        "  datasets = get_dataset(0.125, 256, dataset, subsample_size=1.0)\n",
        "  x_train = datasets[\"x_train\"]\n",
        "  ds_train = datasets[\"ds_train\"]\n",
        "  ds_val = datasets[\"ds_val\"]\n",
        "  ds_test = datasets[\"ds_test\"]\n",
        "  is_classification = datasets[\"is_classification\"]\n",
        "  num_classes = datasets[\"num_classes\"]\n",
        "  num_features = datasets[\"num_features\"]\n",
        "  num_train_steps_select = 1 * len(ds_train) # change back to num_epochs_select\n",
        "  loss_fn = (\n",
        "      tf.keras.losses.CategoricalCrossentropy()\n",
        "      if is_classification\n",
        "      else tf.keras.losses.MeanAbsoluteError()\n",
        "  )\n",
        "\n",
        "  return (\n",
        "      ds_train,\n",
        "      x_train,\n",
        "      ds_val,\n",
        "      ds_test,\n",
        "      is_classification,\n",
        "      num_classes,\n",
        "      num_features,\n",
        "      num_train_steps_select,\n",
        "      loss_fn,\n",
        "      num_epochs_select\n",
        "  )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "arlzXdYQXc0Z"
      },
      "source": [
        "## Experiment utils\n",
        "\n",
        "For both datasets, we ran experiments with a predefined list of feature lengths (20, 30, 40 and 50), to understand the tradeoff between lowering the feature set size and slightly decreasing the model quality.\n",
        "\n",
        "We used a couple of methods to research the differences:\n",
        "  - `SimplSA`: Simple Sequential attention, as in the original paper.\n",
        "  - `SimpleSA20`: SimpleSA, on a 20% random subsample of the dataset.\n",
        "  - `SimpleSA40`: SimpleSA, on a 40% random subsample of the dataset.\n",
        "  - `SimpleSA60`: SimpleSA, on a 60% random subsample of the dataset.\n",
        "  - `IntermediateSA`: Ensembled SimplaSA with K workers, where each worker would provide the next best X features (in our case, 20), and at each checkpoint the ensemble would vote for the best mask. Then, the ensemble will proceed with a common prior, until all feature set sizes are covered.\n",
        "  - `LLYMask`: Feature Selection method according to [Liao et al. 2021](https://arxiv.org/pdf/2010.13631.pdf). Model copied from the Sequential Attention paper code.\n",
        "  - `GroupLasso`: Group lasso feature selection, taken from the original paper code.\n",
        "\n",
        "\n",
        "For each, we run the algorithm 5 times, compute the masks and measure the test set accuracy of a sparse model with a specific mask applied. Also, we measure the time (mean and variance) needed to compute the masks for all 4 feature set sizes.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8r2tdU7hKINf"
      },
      "outputs": [],
      "source": [
        "# Setup\n",
        "feature_counts = [20, 30, 40, 50]\n",
        "random_runs_count = 5\n",
        "methods = [\n",
        "    'SimplSA',\n",
        "    'SimpleSA20',\n",
        "    'SimpleSA40',\n",
        "    'SimpleSA60',\n",
        "    'IntermediateSA',\n",
        "    'LLYMask',\n",
        "    'GroupLasso'\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KanOlNnnK42W"
      },
      "outputs": [],
      "source": [
        "def run_experiment(\n",
        "    get_masks_fn,\n",
        "    get_sparse_model_evaluation,\n",
        "    random_runs: int = random_runs_count,\n",
        "):\n",
        "  random_runs_dict = {}\n",
        "\n",
        "  for i in range(random_runs):\n",
        "    print(f\"Random run {i}/{random_runs}\")\n",
        "    start = time.time()\n",
        "    print(f\"run_experiment dataset size: {len(ds_train)}\")\n",
        "    masks = get_masks_fn()\n",
        "    end = time.time()\n",
        "    print(f\"Masks run {i}: {masks}\")\n",
        "\n",
        "    sparse_model_evaluations = [\n",
        "        get_sparse_model_evaluation(mask) for mask in masks\n",
        "    ]\n",
        "    random_runs_dict[i] = sparse_model_evaluations\n",
        "    random_runs_dict[f\"masks-{i}\"] = masks\n",
        "    random_runs_dict[f\"time-{i}\"] = end - start\n",
        "\n",
        "  return random_runs_dict"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wSZ5bl9xjIcs"
      },
      "outputs": [],
      "source": [
        "def get_data_from_experiment_output(\n",
        "    output, random_runs: int = random_runs_count, divide_time_by: int = 1\n",
        "):\n",
        "  accuracies_for_mask_size = {}\n",
        "  losses_for_mask_size = {}\n",
        "  times = []\n",
        "\n",
        "  for i in range(random_runs):\n",
        "    for feature_count_id in range(len(feature_counts)):\n",
        "      feature_count = feature_counts[feature_count_id]\n",
        "      losses_for_mask_size[feature_count] = losses_for_mask_size.get(\n",
        "          feature_count, []\n",
        "      ) + [output[i][feature_count_id][0]]\n",
        "      accuracies_for_mask_size[feature_count] = accuracies_for_mask_size.get(\n",
        "          feature_count, []\n",
        "      ) + [output[i][feature_count_id][1]]\n",
        "\n",
        "    times.append(output[f\"time-{i}\"] / divide_time_by)\n",
        "\n",
        "  return {\n",
        "      \"mean_acc_for_mask\": {\n",
        "          size: np.mean(accuracies_for_mask_size[size])\n",
        "          for size in accuracies_for_mask_size.keys()\n",
        "      },\n",
        "      \"std_acc_for_mask\": {\n",
        "          size: np.std(accuracies_for_mask_size[size])\n",
        "          for size in accuracies_for_mask_size.keys()\n",
        "      },\n",
        "      \"mean_time\": np.mean(times),\n",
        "      \"std_time\": np.std(times),\n",
        "  }"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zv29UP4pJ41b"
      },
      "source": [
        "## Test run (simple SA)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vBIkaoyjJ73A"
      },
      "source": [
        "\n",
        "\n",
        "## Timed runs (mnist)\n",
        "Goal: track the time it takes to get a mask, then track the performance of a sparse model trained with it"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fmPpdFn1NTDo"
      },
      "outputs": [],
      "source": [
        "# Mnist config.\n",
        "num_epochs_fit = 250"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ug-r3mcCJ3pP"
      },
      "source": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9y3pt2LeNdVQ"
      },
      "source": [
        "### Simple SA"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eaYS38zANiYo"
      },
      "outputs": [],
      "source": [
        "def get_model_simple_sa(to_select):\n",
        "  model = SequentialAttentionMlp(\n",
        "      num_train_steps=num_train_steps_select,\n",
        "      num_classes=num_classes,\n",
        "      layer_sequence=[67],\n",
        "      num_candidates_to_select=to_select,\n",
        "  )\n",
        "\n",
        "  model.compile(\n",
        "      loss=tf.keras.losses.CategoricalCrossentropy(), metrics=[\"accuracy\"]\n",
        "  )\n",
        "\n",
        "  return model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UYcPg4gmTL3k"
      },
      "outputs": [],
      "source": [
        "def get_sparse_model_evaluation_simple_sa(mask):\n",
        "  num_epochs_fit = 250\n",
        "\n",
        "  mlp_fit = SparseModel(\n",
        "      selected_features=tf.math.reduce_sum(\n",
        "          tf.one_hot(mask, num_features, dtype=tf.float32), 0\n",
        "      ),\n",
        "      num_classes=num_classes,\n",
        "      is_classification=True,\n",
        "      layer_sequence=[67],\n",
        "  )\n",
        "  mlp_fit.compile(loss=loss_fn, metrics=[\"accuracy\"])\n",
        "  mlp_fit.fit(\n",
        "      ds_train, validation_data=ds_val, epochs=num_epochs_fit, verbose=None\n",
        "  )\n",
        "  return mlp_fit.evaluate(ds_test)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dNHB-jOXSB5R"
      },
      "outputs": [],
      "source": [
        "def get_masks_fn_simple_sa():\n",
        "  masks = []\n",
        "  for feature_count in feature_counts:\n",
        "    print(f\"\\_ Getting mask for {feature_count} features.\")\n",
        "    mask = []\n",
        "\n",
        "    model = get_model_simple_sa(feature_count)\n",
        "    model.fit(\n",
        "        ds_train, validation_data=ds_val, epochs=num_epochs_select, verbose=None\n",
        "    )\n",
        "\n",
        "    selected_features = model.sa.selected_features\n",
        "    _, selected_indices = tf.math.top_k(selected_features, k=feature_count)\n",
        "    selected_indices = selected_indices.numpy()\n",
        "\n",
        "    masks.append(list(selected_indices))\n",
        "\n",
        "  return masks"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sYdP4jSfVwUw"
      },
      "outputs": [],
      "source": [
        "# Reload global variables.\n",
        "(\n",
        "    ds_train,\n",
        "    x_train,\n",
        "    ds_val,\n",
        "    ds_test,\n",
        "    is_classification,\n",
        "    num_classes,\n",
        "    num_features,\n",
        "    num_train_steps_select,\n",
        "    loss_fn,\n",
        "    num_epochs_select,\n",
        ") = load_dataset(num_epochs_select=50, batch_size=256, dataset='mnist')\n",
        "simple_sa_experiment = run_experiment(\n",
        "    get_masks_fn_simple_sa, get_sparse_model_evaluation_simple_sa\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3UKXb7wFx6ma"
      },
      "outputs": [],
      "source": [
        "print(simple_sa_experiment)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "moNqfDANiYi8"
      },
      "outputs": [],
      "source": [
        "simple_sa_experiment_data = {0: [[0.8113244771957397, 0.7653999924659729], [0.6454375982284546, 0.8220000267028809], [0.5484784245491028, 0.8528000116348267], [0.5065304040908813, 0.8694999814033508]], 'masks-0': [[101, 152, 212, 220, 265, 289, 319, 376, 378, 401, 426, 431, 461, 464, 518, 543, 595, 597, 655, 714], [98, 103, 152, 183, 208, 212, 219, 264, 267, 280, 325, 347, 349, 350, 356, 373, 375, 410, 431, 434, 454, 457, 460, 514, 545, 571, 579, 596, 657, 711], [1, 102, 150, 154, 180, 184, 189, 201, 210, 221, 233, 239, 241, 248, 264, 266, 271, 273, 316, 320, 322, 324, 346, 374, 378, 385, 400, 434, 455, 464, 485, 488, 489, 491, 539, 542, 573, 596, 657, 658], [106, 128, 152, 163, 175, 185, 206, 209, 211, 220, 241, 247, 268, 273, 288, 295, 298, 304, 318, 319, 341, 343, 353, 354, 375, 376, 405, 409, 428, 430, 435, 442, 462, 466, 485, 514, 515, 537, 549, 555, 568, 573, 598, 607, 650, 652, 658, 707, 712, 717]], 'time-0': 7.725859642028809}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qGuK5yzKZHzE"
      },
      "source": [
        "### Simple SA 20\n",
        "Sequential attention on 20% of the dataset (randomly sampled)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BS2mVZXZZTPk"
      },
      "outputs": [],
      "source": [
        "def get_model_simple_sa20(to_select):\n",
        "  return get_model_simple_sa(to_select)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WcD3JfnYZlYY"
      },
      "outputs": [],
      "source": [
        "def get_sparse_model_evaluation_simple_sa20(mask):\n",
        "  return get_sparse_model_evaluation_simple_sa(mask)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YkmzFn32ZqA3"
      },
      "outputs": [],
      "source": [
        "def get_masks_fn_simple_sa20():\n",
        "  return get_masks_fn_simple_sa()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xt7Lp6XhZqKY"
      },
      "outputs": [],
      "source": [
        "# Reload global variables.\n",
        "(\n",
        "    ds_train,\n",
        "    x_train,\n",
        "    ds_val,\n",
        "    ds_test,\n",
        "    is_classification,\n",
        "    num_classes,\n",
        "    num_features,\n",
        "    num_train_steps_select,\n",
        "    loss_fn,\n",
        "    num_epochs_select,\n",
        ") = load_dataset(\n",
        "    num_epochs_select=50, batch_size=256, dataset='mnist', subsample_size=0.2\n",
        ")\n",
        "simple_sa_20_experiment_data = run_experiment(\n",
        "    get_masks_fn_simple_sa20, get_sparse_model_evaluation_simple_sa20\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dP0bqpzI9-OD"
      },
      "outputs": [],
      "source": [
        "print(simple_sa_20_experiment_data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6Pj2t7OYl38x"
      },
      "outputs": [],
      "source": [
        "simple_sa_20_experiment_data = {0: [[0.47898370027542114, 0.8636000156402588], [0.4648936986923218, 0.9035000205039978], [0.5981168746948242, 0.9150000214576721], [0.605646550655365, 0.9266999959945679]], 'masks-0': [[153, 211, 248, 267, 269, 290, 326, 373, 375, 378, 409, 427, 433, 462, 478, 484, 514, 544, 568, 657], [138, 152, 156, 211, 213, 248, 267, 270, 317, 321, 323, 326, 347, 357, 378, 381, 399, 403, 429, 434, 436, 482, 486, 490, 516, 568, 570, 572, 580, 657], [45, 122, 153, 157, 191, 203, 210, 240, 248, 267, 270, 291, 294, 300, 323, 330, 345, 348, 354, 375, 378, 380, 399, 401, 409, 434, 454, 459, 464, 467, 485, 490, 516, 542, 574, 594, 598, 656, 661, 708], [152, 161, 184, 206, 210, 212, 220, 237, 242, 267, 277, 290, 293, 297, 316, 322, 323, 327, 330, 347, 373, 376, 378, 381, 399, 404, 430, 434, 439, 464, 483, 485, 487, 489, 526, 542, 544, 546, 566, 568, 576, 578, 580, 599, 610, 650, 657, 709, 720, 752]], 'time-0': 27.11506223678589, 1: [[0.45385250449180603, 0.8748999834060669], [0.5053437352180481, 0.9013000130653381], [0.5983988046646118, 0.9160000085830688], [0.6262752413749695, 0.9271000027656555]], 'masks-1': [[153, 210, 213, 267, 296, 326, 345, 375, 377, 379, 433, 436, 455, 459, 485, 490, 516, 542, 567, 657], [152, 157, 203, 210, 212, 242, 248, 267, 291, 322, 324, 326, 343, 345, 347, 359, 378, 384, 403, 409, 427, 434, 458, 488, 490, 542, 550, 567, 572, 657], [63, 152, 156, 158, 207, 210, 213, 219, 244, 265, 267, 276, 290, 296, 298, 322, 343, 347, 356, 373, 376, 378, 381, 387, 427, 430, 434, 436, 439, 489, 513, 515, 518, 554, 571, 595, 601, 657, 708, 710], [126, 151, 156, 182, 191, 203, 209, 214, 238, 240, 247, 267, 270, 277, 290, 301, 321, 323, 326, 343, 345, 347, 356, 359, 375, 378, 380, 401, 409, 427, 432, 434, 465, 484, 486, 488, 490, 503, 541, 543, 550, 564, 568, 600, 610, 656, 659, 708, 710, 719]], 'time-1': 28.521087884902954, 2: [[0.4909895956516266, 0.8641999959945679], [0.48868364095687866, 0.8991000056266785], [0.5978553891181946, 0.9140999913215637], [0.567193329334259, 0.9319999814033508]], 'masks-2': [[152, 210, 220, 268, 298, 317, 347, 378, 381, 401, 427, 434, 459, 465, 490, 496, 513, 543, 567, 656], [153, 176, 186, 191, 209, 211, 248, 267, 269, 317, 322, 324, 327, 347, 371, 377, 381, 387, 428, 430, 434, 436, 490, 514, 544, 567, 570, 574, 657, 710], [153, 156, 175, 191, 210, 212, 248, 267, 270, 290, 302, 306, 320, 323, 344, 350, 354, 375, 380, 399, 401, 406, 436, 439, 459, 462, 482, 485, 488, 518, 543, 570, 573, 578, 583, 595, 656, 687, 708, 719], [153, 158, 177, 184, 191, 208, 210, 212, 214, 266, 267, 269, 276, 291, 298, 300, 317, 321, 323, 331, 344, 348, 354, 357, 375, 378, 380, 399, 401, 404, 408, 433, 439, 454, 456, 458, 462, 465, 513, 516, 518, 542, 567, 572, 581, 656, 661, 708, 710, 719]], 'time-2': 30.152465105056763, 3: [[0.4769500494003296, 0.871399998664856], [0.5713586211204529, 0.8934999704360962], [0.5689853429794312, 0.9157999753952026], [0.6094802021980286, 0.9287999868392944]], 'masks-3': [[144, 153, 184, 210, 248, 267, 270, 319, 345, 353, 376, 406, 427, 436, 458, 488, 490, 542, 597, 656], [153, 157, 175, 209, 211, 241, 248, 266, 295, 297, 317, 326, 328, 347, 359, 378, 403, 428, 434, 436, 486, 490, 493, 516, 569, 571, 594, 655, 657, 708], [152, 157, 175, 184, 191, 210, 240, 242, 248, 267, 297, 317, 319, 322, 349, 354, 358, 371, 375, 378, 380, 401, 427, 433, 436, 458, 462, 467, 484, 516, 542, 546, 568, 572, 580, 594, 656, 661, 708, 720], [20, 153, 156, 174, 177, 191, 210, 212, 216, 232, 241, 248, 263, 265, 267, 270, 302, 315, 321, 323, 328, 345, 347, 353, 371, 378, 401, 404, 409, 414, 427, 434, 458, 464, 467, 488, 490, 512, 516, 542, 567, 572, 580, 583, 597, 656, 660, 708, 710, 719]], 'time-3': 28.977137565612793, 4: [[0.48589780926704407, 0.8632000088691711], [0.4646494686603546, 0.9057000279426575], [0.542389452457428, 0.9172999858856201], [0.5538588166236877, 0.9337999820709229]], 'masks-4': [[84, 153, 210, 240, 248, 269, 290, 323, 326, 374, 377, 399, 406, 436, 458, 490, 516, 539, 542, 657], [152, 156, 159, 191, 203, 210, 212, 214, 248, 264, 267, 269, 273, 326, 343, 345, 347, 350, 405, 409, 427, 458, 462, 488, 541, 543, 550, 567, 574, 657], [152, 174, 184, 190, 210, 241, 248, 267, 271, 290, 292, 296, 322, 326, 344, 348, 357, 374, 377, 380, 387, 427, 430, 433, 437, 463, 467, 484, 486, 516, 517, 542, 567, 574, 599, 610, 656, 708, 710, 719], [99, 106, 149, 153, 157, 184, 207, 210, 216, 240, 248, 265, 267, 270, 291, 294, 296, 316, 326, 328, 330, 345, 347, 350, 371, 376, 380, 381, 401, 403, 405, 426, 434, 436, 439, 456, 460, 486, 489, 497, 518, 544, 550, 567, 570, 583, 600, 656, 708, 719]], 'time-4': 32.11427402496338}\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "esQD5uR3ad3Q"
      },
      "source": [
        "### Simple SA 40\n",
        "Sequential attention on 40% of the dataset (randomly sampled)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mLsPgQp3ahal"
      },
      "outputs": [],
      "source": [
        "def get_model_simple_sa40(to_select):\n",
        "  return get_model_simple_sa(to_select)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ugq5iLYOahal"
      },
      "outputs": [],
      "source": [
        "def get_sparse_model_evaluation_simple_sa40(mask):\n",
        "  return get_sparse_model_evaluation_simple_sa(mask)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LqxTDOsPahal"
      },
      "outputs": [],
      "source": [
        "def get_masks_fn_simple_sa40():\n",
        "  return get_masks_fn_simple_sa()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AMq28F58aoZy"
      },
      "outputs": [],
      "source": [
        "# Reload global variables.\n",
        "(\n",
        "    ds_train,\n",
        "    x_train,\n",
        "    ds_val,\n",
        "    ds_test,\n",
        "    is_classification,\n",
        "    num_classes,\n",
        "    num_features,\n",
        "    num_train_steps_select,\n",
        "    loss_fn,\n",
        "    num_epochs_select,\n",
        ") = load_dataset(\n",
        "    num_epochs_select=50, batch_size=256, dataset='mnist', subsample_size=0.4\n",
        ")\n",
        "simple_sa_40_experiment_data = run_experiment(\n",
        "    get_masks_fn_simple_sa40, get_sparse_model_evaluation_simple_sa40\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XEDrUp5T-C2O"
      },
      "outputs": [],
      "source": [
        "print(simple_sa_40_experiment_data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dWliC_V_l_dm"
      },
      "outputs": [],
      "source": [
        "simple_sa_40_experiment_data = {0: [[0.3951590359210968, 0.883899986743927], [0.35931286215782166, 0.9049999713897705], [0.4721135199069977, 0.921999990940094], [0.5569542050361633, 0.932699978351593]], 'masks-0': [[150, 156, 181, 211, 248, 267, 269, 299, 318, 350, 375, 406, 409, 428, 458, 489, 539, 542, 572, 627], [7, 150, 156, 181, 191, 211, 239, 248, 290, 296, 298, 315, 321, 328, 346, 375, 378, 386, 397, 400, 409, 433, 458, 462, 539, 542, 544, 580, 627, 657], [125, 156, 177, 181, 191, 211, 236, 239, 246, 248, 268, 270, 272, 289, 318, 321, 347, 350, 352, 354, 376, 386, 398, 400, 402, 406, 409, 455, 457, 461, 464, 495, 514, 518, 544, 567, 570, 575, 583, 657], [99, 124, 155, 177, 181, 186, 191, 211, 235, 238, 248, 261, 264, 266, 269, 290, 296, 299, 330, 343, 348, 350, 354, 372, 374, 380, 387, 404, 425, 428, 434, 437, 459, 463, 483, 486, 488, 495, 513, 517, 542, 555, 572, 575, 580, 596, 631, 657, 708, 711]], 'time-0': 46.07294273376465, 1: [[0.39046767354011536, 0.8840000033378601], [0.35353150963783264, 0.914900004863739], [0.462999165058136, 0.9302999973297119], [0.5590104460716248, 0.9365000128746033]], 'masks-1': [[151, 182, 239, 248, 290, 296, 298, 328, 350, 374, 404, 409, 427, 434, 457, 517, 542, 567, 572, 628], [99, 151, 156, 176, 181, 191, 211, 236, 248, 266, 269, 317, 320, 326, 328, 350, 371, 374, 386, 406, 428, 437, 461, 486, 490, 539, 542, 544, 627, 657], [124, 156, 177, 181, 188, 213, 219, 237, 239, 261, 267, 271, 276, 297, 300, 318, 320, 343, 350, 354, 359, 376, 380, 397, 401, 404, 427, 434, 437, 484, 486, 488, 490, 495, 544, 567, 570, 580, 656, 659], [125, 156, 177, 181, 191, 211, 235, 238, 242, 248, 266, 268, 299, 301, 318, 320, 321, 325, 343, 345, 350, 354, 359, 375, 377, 379, 397, 400, 402, 404, 409, 411, 434, 455, 458, 465, 488, 490, 513, 539, 542, 544, 550, 565, 575, 580, 627, 657, 707, 710]], 'time-1': 48.18356490135193, 2: [[0.3833909034729004, 0.8873999714851379], [0.35052865743637085, 0.9140999913215637], [0.43153491616249084, 0.9222999811172485], [0.5783820748329163, 0.9337999820709229]], 'masks-2': [[150, 156, 181, 211, 267, 297, 300, 317, 350, 359, 374, 404, 409, 434, 456, 486, 516, 570, 595, 657], [101, 124, 156, 177, 183, 236, 239, 248, 270, 272, 289, 296, 326, 346, 348, 350, 387, 397, 400, 404, 406, 409, 456, 459, 462, 514, 523, 539, 544, 626], [99, 124, 156, 177, 181, 191, 211, 215, 236, 239, 247, 268, 298, 304, 316, 318, 321, 328, 350, 375, 381, 387, 397, 400, 402, 406, 437, 456, 461, 463, 486, 517, 539, 542, 544, 555, 565, 626, 657, 708], [99, 124, 148, 156, 178, 181, 183, 210, 213, 219, 233, 236, 239, 266, 269, 271, 276, 295, 318, 321, 328, 343, 348, 350, 352, 354, 358, 374, 397, 400, 404, 406, 409, 431, 439, 456, 462, 486, 488, 518, 523, 542, 544, 554, 567, 600, 626, 657, 661, 708]], 'time-2': 51.737651109695435, 3: [[0.3820122182369232, 0.883899986743927], [0.3405912220478058, 0.9089999794960022], [0.44281700253486633, 0.9254999756813049], [0.5184953212738037, 0.9390000104904175]], 'masks-3': [[151, 156, 182, 211, 267, 298, 318, 348, 358, 371, 378, 402, 409, 456, 461, 514, 523, 544, 596, 657], [5, 125, 156, 177, 181, 211, 238, 241, 248, 272, 290, 295, 298, 320, 352, 371, 374, 377, 387, 434, 437, 456, 486, 488, 491, 517, 542, 583, 596, 657], [89, 122, 125, 156, 177, 181, 190, 211, 213, 248, 267, 270, 290, 296, 299, 321, 329, 347, 350, 353, 371, 373, 376, 386, 397, 406, 431, 437, 456, 461, 463, 467, 486, 518, 539, 542, 544, 582, 596, 657], [124, 127, 177, 181, 191, 211, 213, 234, 237, 248, 267, 269, 271, 293, 305, 318, 324, 328, 343, 348, 350, 354, 358, 373, 375, 380, 398, 406, 431, 437, 439, 456, 460, 463, 486, 489, 511, 523, 544, 564, 570, 574, 577, 581, 596, 654, 657, 708, 710, 715]], 'time-3': 51.56308889389038, 4: [[0.39909055829048157, 0.8806999921798706], [0.3723650276660919, 0.90829998254776], [0.42380985617637634, 0.925000011920929], [0.5662999749183655, 0.9326000213623047]], 'masks-4': [[150, 181, 238, 248, 269, 290, 299, 348, 350, 371, 374, 404, 409, 434, 456, 488, 542, 545, 596, 656], [125, 156, 177, 183, 191, 236, 239, 248, 269, 299, 317, 321, 343, 347, 350, 353, 386, 401, 406, 409, 455, 461, 486, 491, 544, 570, 580, 596, 657, 749], [124, 155, 176, 181, 191, 211, 237, 248, 266, 268, 272, 290, 298, 320, 324, 343, 350, 354, 374, 376, 380, 386, 400, 403, 434, 437, 455, 459, 463, 488, 513, 517, 523, 539, 543, 556, 595, 600, 657, 708], [122, 125, 156, 176, 178, 183, 191, 208, 212, 215, 239, 248, 261, 264, 268, 293, 298, 301, 315, 316, 318, 324, 348, 350, 354, 380, 387, 397, 401, 403, 406, 409, 427, 438, 456, 460, 463, 486, 489, 523, 541, 544, 546, 555, 567, 570, 628, 657, 708, 722]], 'time-4': 48.25512480735779}\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nUokMIusYsDe"
      },
      "source": [
        "### Simple SA 60\n",
        "Sequential attention on 60% of the dataset (randomly sampled)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tIy4qsTQYsD5"
      },
      "outputs": [],
      "source": [
        "def get_model_simple_sa60(to_select):\n",
        "  return get_model_simple_sa(to_select)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ioxi4V__YsD6"
      },
      "outputs": [],
      "source": [
        "def get_sparse_model_evaluation_simple_sa60(mask):\n",
        "  return get_sparse_model_evaluation_simple_sa(mask)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SWUQ2srrYsD6"
      },
      "outputs": [],
      "source": [
        "def get_masks_fn_simple_sa60():\n",
        "  return get_masks_fn_simple_sa()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jjyGBWh6YsD7"
      },
      "outputs": [],
      "source": [
        "# Reload global variables.\n",
        "(\n",
        "    ds_train,\n",
        "    x_train,\n",
        "    ds_val,\n",
        "    ds_test,\n",
        "    is_classification,\n",
        "    num_classes,\n",
        "    num_features,\n",
        "    num_train_steps_select,\n",
        "    loss_fn,\n",
        "    num_epochs_select,\n",
        ") = load_dataset(\n",
        "    num_epochs_select=50, batch_size=256, dataset='mnist', subsample_size=0.6\n",
        ")\n",
        "simple_sa_60_experiment_data = run_experiment(\n",
        "    get_masks_fn_simple_sa60, get_sparse_model_evaluation_simple_sa60\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NptalamuYsD7"
      },
      "outputs": [],
      "source": [
        "print(simple_sa_60_experiment_data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xZs-bs1eYsD8"
      },
      "outputs": [],
      "source": [
        "simple_sa_60_experiment_data = {0: [[0.3700980544090271, 0.8835999965667725], [0.2911509871482849, 0.9229000210762024], [0.3522627353668213, 0.9336000084877014], [0.41702744364738464, 0.9413999915122986]], 'masks-0': [[151, 155, 211, 220, 263, 269, 317, 326, 350, 374, 376, 409, 427, 434, 458, 517, 522, 542, 567, 657], [31, 153, 156, 177, 211, 248, 263, 267, 269, 300, 317, 319, 326, 350, 352, 358, 371, 373, 376, 409, 430, 434, 455, 486, 489, 518, 543, 569, 595, 656], [103, 152, 155, 177, 181, 187, 212, 235, 238, 247, 269, 292, 294, 298, 316, 318, 328, 342, 350, 358, 374, 376, 380, 399, 406, 409, 430, 455, 461, 464, 468, 485, 487, 518, 528, 543, 565, 568, 580, 657], [126, 151, 156, 181, 192, 210, 212, 215, 231, 238, 248, 263, 269, 293, 295, 299, 301, 316, 318, 324, 348, 350, 358, 372, 374, 378, 381, 399, 405, 431, 436, 439, 455, 461, 465, 485, 487, 518, 538, 541, 543, 545, 568, 580, 602, 655, 658, 712, 719, 745]], 'time-0': 69.08614659309387, 1: [[0.36893439292907715, 0.8862000107765198], [0.29431915283203125, 0.9215999841690063], [0.3281683921813965, 0.9340999722480774], [0.43689867854118347, 0.9422000050544739]], 'masks-1': [[45, 150, 154, 211, 266, 297, 299, 317, 348, 350, 374, 386, 409, 433, 456, 489, 541, 543, 566, 655], [66, 150, 154, 211, 219, 263, 266, 269, 272, 317, 326, 346, 350, 358, 371, 376, 409, 431, 434, 455, 457, 487, 490, 513, 543, 545, 568, 580, 655, 658], [27, 150, 154, 209, 211, 214, 219, 263, 266, 269, 272, 276, 295, 317, 320, 326, 346, 350, 352, 358, 371, 376, 409, 425, 431, 434, 456, 464, 486, 488, 510, 513, 518, 543, 568, 572, 581, 655, 658, 712], [97, 151, 154, 157, 176, 209, 212, 215, 219, 238, 261, 263, 269, 276, 295, 299, 301, 317, 326, 348, 349, 352, 357, 372, 374, 378, 387, 397, 399, 404, 409, 431, 434, 455, 457, 464, 487, 489, 495, 514, 518, 540, 543, 554, 572, 595, 603, 606, 655, 658]], 'time-1': 76.09812188148499, 2: [[0.3822028636932373, 0.8810999989509583], [0.28970280289649963, 0.920199990272522], [0.32942426204681396, 0.9325000047683716], [0.43238911032676697, 0.9386000037193298]], 'masks-2': [[154, 177, 211, 219, 269, 299, 346, 348, 350, 371, 405, 409, 456, 461, 486, 518, 543, 569, 656, 760], [98, 155, 177, 211, 248, 263, 266, 269, 272, 295, 317, 326, 348, 350, 358, 374, 399, 405, 409, 434, 456, 459, 486, 517, 539, 542, 550, 597, 655, 687], [6, 123, 154, 157, 177, 209, 212, 214, 219, 244, 264, 267, 276, 289, 296, 298, 328, 350, 374, 376, 380, 386, 399, 405, 409, 429, 434, 455, 464, 485, 487, 490, 541, 543, 545, 554, 565, 568, 656, 660], [137, 151, 154, 156, 176, 209, 211, 215, 220, 260, 263, 266, 269, 276, 292, 295, 317, 324, 326, 328, 343, 347, 350, 376, 386, 399, 401, 405, 408, 425, 434, 437, 455, 459, 484, 486, 490, 501, 516, 522, 542, 545, 554, 565, 568, 572, 580, 655, 657, 660]], 'time-2': 70.31428408622742, 3: [[0.37265118956565857, 0.8859000205993652], [0.27209287881851196, 0.9251000285148621], [0.34308189153671265, 0.934499979019165], [0.42168310284614563, 0.9401999711990356]], 'masks-3': [[35, 150, 154, 211, 263, 268, 298, 330, 347, 350, 373, 376, 409, 427, 434, 485, 518, 543, 568, 656], [123, 155, 177, 211, 235, 241, 248, 266, 297, 299, 319, 323, 330, 345, 350, 376, 381, 398, 402, 406, 456, 461, 464, 490, 514, 523, 544, 568, 655, 658], [126, 151, 155, 177, 188, 209, 211, 248, 261, 266, 268, 298, 300, 304, 317, 319, 343, 348, 350, 352, 357, 374, 387, 399, 405, 409, 434, 456, 459, 464, 486, 490, 513, 543, 545, 565, 568, 580, 655, 658], [124, 156, 181, 187, 202, 210, 212, 234, 242, 248, 263, 267, 293, 296, 316, 325, 327, 346, 348, 350, 358, 399, 401, 404, 406, 409, 415, 425, 440, 456, 458, 462, 487, 512, 518, 523, 539, 542, 544, 554, 576, 580, 599, 625, 656, 658, 708, 711, 719, 742]], 'time-3': 68.37718534469604, 4: [[0.38910746574401855, 0.8813999891281128], [0.28183579444885254, 0.9228000044822693], [0.34288230538368225, 0.9344000220298767], [0.4118873178958893, 0.9440000057220459]], 'masks-4': [[7, 151, 184, 238, 263, 269, 326, 344, 350, 374, 405, 409, 456, 461, 486, 518, 538, 543, 569, 655], [97, 150, 154, 204, 211, 248, 266, 269, 291, 299, 317, 348, 350, 353, 358, 374, 399, 405, 409, 429, 461, 485, 514, 518, 539, 544, 570, 596, 656, 770], [151, 155, 181, 184, 191, 212, 238, 248, 263, 269, 295, 300, 317, 326, 347, 349, 352, 358, 376, 378, 397, 399, 401, 409, 432, 461, 464, 467, 484, 486, 510, 518, 543, 545, 569, 579, 595, 602, 656, 659], [122, 153, 157, 183, 204, 209, 211, 214, 219, 234, 238, 240, 248, 263, 266, 296, 298, 300, 303, 317, 321, 347, 350, 352, 355, 358, 370, 372, 374, 376, 405, 409, 429, 434, 455, 459, 464, 468, 489, 513, 518, 526, 543, 567, 569, 573, 580, 629, 655, 659]], 'time-4': 74.30822348594666}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HMFSdQ2da6se"
      },
      "source": [
        "### Intermediate SA\n",
        "In comparison to Simple SA, Intermediate SA has a pool of K workers, that each compute the next best N features. At each step, the ensemble votes for a common mask, and all proceed with the same prior. In a setup where we need masks for [10, 20, 30] features, the ensemble first votes for 10 features, then with the common prior the ensemble votes again for the next 10 features, and so on until completion."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8yF3cXKua-pA"
      },
      "outputs": [],
      "source": [
        "def get_model_intermediate_sa(\n",
        "    preselected_indices, to_select, num_train_steps_select\n",
        "):\n",
        "  model = SequentialAttentionMlp(\n",
        "      num_train_steps=num_train_steps_select,\n",
        "      num_classes=num_classes,\n",
        "      layer_sequence=[67],\n",
        "      num_candidates_to_select=to_select,\n",
        "      preselected_candidate_indices=preselected_indices,\n",
        "  )\n",
        "\n",
        "  model.compile(\n",
        "      loss=tf.keras.losses.CategoricalCrossentropy(), metrics=[\"accuracy\"]\n",
        "  )\n",
        "\n",
        "  return model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zyQ0kCJ_a-pB"
      },
      "outputs": [],
      "source": [
        "def get_sparse_model_evaluation_intermediate_sa(mask):\n",
        "  mlp_fit = SparseModel(\n",
        "      selected_features=tf.math.reduce_sum(\n",
        "          tf.one_hot(mask, num_features, dtype=tf.float32), 0\n",
        "      ),\n",
        "      num_classes=num_classes,\n",
        "      is_classification=True,\n",
        "      layer_sequence=[67],\n",
        "  )\n",
        "  mlp_fit.compile(loss=loss_fn, metrics=[\"accuracy\"])\n",
        "  mlp_fit.fit(ds_train, validation_data=ds_val, epochs=250, verbose=None)\n",
        "  return mlp_fit.evaluate(ds_test)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8OIe-i4Ta-pB"
      },
      "outputs": [],
      "source": [
        "def get_masks_fn_intermediate_sa():\n",
        "  # parameters adapted for the constants defined in feature_counts\n",
        "  pool_size = 5\n",
        "  num_features_at_each_step = 30  # the gap for the features is at most 20; so at each step, we aim to fill preselected_features to the desired capacity\n",
        "  epochs_for_each = num_epochs_select\n",
        "\n",
        "  preselected_features = []\n",
        "\n",
        "  for feature_target in feature_counts:\n",
        "    print(f\"\\_ Getting mask for {feature_target} features.\")\n",
        "    model_pool = [\n",
        "        get_model_intermediate_sa(\n",
        "            preselected_features,\n",
        "            feature_target,\n",
        "            epochs_for_each * len(ds_train),\n",
        "        )\n",
        "        for i in range(pool_size)\n",
        "    ]\n",
        "\n",
        "    [\n",
        "        model.fit(\n",
        "            ds_train,\n",
        "            validation_data=ds_val,\n",
        "            epochs=epochs_for_each,\n",
        "            verbose=None,\n",
        "        )\n",
        "        for model in model_pool\n",
        "    ]\n",
        "\n",
        "    votes = {}\n",
        "\n",
        "    for model in model_pool:\n",
        "      selected_features = model.sa.selected_features\n",
        "      _, selected_indices = tf.math.top_k(selected_features, k=feature_target)\n",
        "      selected_indices = selected_indices.numpy()\n",
        "      for feature in list(selected_indices):\n",
        "        if feature not in preselected_features:\n",
        "          votes[feature] = votes.get(feature, 0) + 1\n",
        "\n",
        "    while len(preselected_features) \u003c feature_target:\n",
        "      max_feature = max(votes, key=votes.get)\n",
        "      votes[max_feature] = 0\n",
        "      preselected_features.append(max_feature)\n",
        "\n",
        "    print(f\"Features: {preselected_features}\")\n",
        "\n",
        "  return [\n",
        "      preselected_features[: feature_target + 1]\n",
        "      for feature_target in feature_counts\n",
        "  ]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JKANZ7A6HAN5"
      },
      "outputs": [],
      "source": [
        "# Reload global variables.\n",
        "(\n",
        "    ds_train,\n",
        "    x_train,\n",
        "    ds_val,\n",
        "    ds_test,\n",
        "    is_classification,\n",
        "    num_classes,\n",
        "    num_features,\n",
        "    num_train_steps_select,\n",
        "    loss_fn,\n",
        "    num_epochs_select,\n",
        ") = load_dataset(num_epochs_select=50, batch_size=256, dataset='mnist')\n",
        "intermediate_sa_experiment_data = run_experiment(\n",
        "    get_masks_fn_intermediate_sa, get_sparse_model_evaluation_intermediate_sa\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "94gYcvQVbcEj"
      },
      "outputs": [],
      "source": [
        "print(intermediate_sa_experiment_data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9MtwOagPmH0T"
      },
      "outputs": [],
      "source": [
        "intermediate_sa_experiment_data = {0: [[0.36737897992134094, 0.8870999813079834], [0.25769931077957153, 0.9261000156402588], [0.22425919771194458, 0.9427000284194946], [0.24101532995700836, 0.9513000249862671]], 'masks-0': [[326, 409, 239, 269, 378, 461, 657, 490, 514, 151, 155, 290, 401, 595, 386, 572, 154, 211, 375, 350, 276], [326, 409, 239, 269, 378, 461, 657, 490, 514, 151, 155, 290, 401, 595, 386, 572, 154, 211, 375, 350, 276, 320, 328, 432, 464, 555, 220, 294, 343, 404, 352], [326, 409, 239, 269, 378, 461, 657, 490, 514, 151, 155, 290, 401, 595, 386, 572, 154, 211, 375, 350, 276, 320, 328, 432, 464, 555, 220, 294, 343, 404, 352, 209, 345, 570, 243, 457, 484, 208, 455, 516, 369], [326, 409, 239, 269, 378, 461, 657, 490, 514, 151, 155, 290, 401, 595, 386, 572, 154, 211, 375, 350, 276, 320, 328, 432, 464, 555, 220, 294, 343, 404, 352, 209, 345, 570, 243, 457, 484, 208, 455, 516, 369, 97, 473, 579, 202, 407, 439, 708, 161, 388]], 'time-0': 140.84598684310913, 1: [[0.39949676394462585, 0.8755000233650208], [0.2562074661254883, 0.925599992275238], [0.23446759581565857, 0.9409000277519226], [0.2668888568878174, 0.9424999952316284]], 'masks-1': [[220, 569, 267, 409, 154, 211, 350, 375, 406, 514, 566, 656, 150, 212, 484, 490, 655, 153, 210, 318, 544], [220, 569, 267, 409, 154, 211, 350, 375, 406, 514, 566, 656, 150, 212, 484, 490, 655, 153, 210, 318, 544, 276, 460, 298, 659, 328, 400, 277, 296, 371, 358], [220, 569, 267, 409, 154, 211, 350, 375, 406, 514, 566, 656, 150, 212, 484, 490, 655, 153, 210, 318, 544, 276, 460, 298, 659, 328, 400, 277, 296, 371, 358, 158, 202, 430, 157, 320, 380, 744, 348, 397, 316], [220, 569, 267, 409, 154, 211, 350, 375, 406, 514, 566, 656, 150, 212, 484, 490, 655, 153, 210, 318, 544, 276, 460, 298, 659, 328, 400, 277, 296, 371, 358, 158, 202, 430, 157, 320, 380, 744, 348, 397, 316, 235, 464, 528, 720, 208, 439, 218, 232, 501]], 'time-1': 122.75689029693604, 2: [[0.36665841937065125, 0.883400022983551], [0.24630942940711975, 0.9290000200271606], [0.22387275099754333, 0.9380000233650208], [0.25911858677864075, 0.9467999935150146]], 'masks-2': [[150, 248, 406, 239, 347, 350, 656, 155, 211, 376, 458, 514, 539, 571, 657, 210, 297, 345, 428, 462, 191], [150, 248, 406, 239, 347, 350, 656, 155, 211, 376, 458, 514, 539, 571, 657, 210, 297, 345, 428, 462, 191, 516, 409, 518, 398, 291, 300, 243, 268, 327, 305], [150, 248, 406, 239, 347, 350, 656, 155, 211, 376, 458, 514, 539, 571, 657, 210, 297, 345, 428, 462, 191, 516, 409, 518, 398, 291, 300, 243, 268, 327, 305, 352, 402, 555, 660, 233, 578, 208, 315, 343, 213], [150, 248, 406, 239, 347, 350, 656, 155, 211, 376, 458, 514, 539, 571, 657, 210, 297, 345, 428, 462, 191, 516, 409, 518, 398, 291, 300, 243, 268, 327, 305, 352, 402, 555, 660, 233, 578, 208, 315, 343, 213, 321, 464, 203, 466, 575, 602, 330, 202, 329]], 'time-2': 124.29457783699036, 3: [[0.358131468296051, 0.8867999911308289], [0.25287866592407227, 0.9247999787330627], [0.22342118620872498, 0.9395999908447266], [0.23133215308189392, 0.9473000168800354]], 'masks-3': [[211, 268, 347, 378, 154, 248, 298, 348, 358, 399, 429, 462, 657, 381, 461, 513, 596, 326, 403, 464, 317], [211, 268, 347, 378, 154, 248, 298, 348, 358, 399, 429, 462, 657, 381, 461, 513, 596, 326, 403, 464, 317, 328, 176, 322, 487, 96, 528, 97, 190, 237, 438], [211, 268, 347, 378, 154, 248, 298, 348, 358, 399, 429, 462, 657, 381, 461, 513, 596, 326, 403, 464, 317, 328, 176, 322, 487, 96, 528, 97, 190, 237, 438, 571, 266, 543, 157, 352, 207, 516, 660, 156, 373], [211, 268, 347, 378, 154, 248, 298, 348, 358, 399, 429, 462, 657, 381, 461, 513, 596, 326, 403, 464, 317, 328, 176, 322, 487, 96, 528, 97, 190, 237, 438, 571, 266, 543, 157, 352, 207, 516, 660, 156, 373, 433, 518, 679, 214, 405, 511, 708, 719, 213]], 'time-3': 123.7198314666748, 4: [[0.3678957521915436, 0.8884999752044678], [0.2627558708190918, 0.9251999855041504], [0.24510514736175537, 0.939300000667572], [0.2638625204563141, 0.9448000192642212]], 'masks-4': [[347, 543, 211, 269, 378, 399, 248, 656, 150, 267, 436, 568, 657, 183, 239, 317, 381, 513, 156, 462, 299], [347, 543, 211, 269, 378, 399, 248, 656, 150, 267, 436, 568, 657, 183, 239, 317, 381, 513, 156, 462, 299, 321, 430, 488, 191, 404, 358, 401, 431, 517, 208], [347, 543, 211, 269, 378, 399, 248, 656, 150, 267, 436, 568, 657, 183, 239, 317, 381, 513, 156, 462, 299, 321, 430, 488, 191, 404, 358, 401, 431, 517, 208, 202, 661, 153, 291, 355, 455, 555, 528, 593, 305], [347, 543, 211, 269, 378, 399, 248, 656, 150, 267, 436, 568, 657, 183, 239, 317, 381, 513, 156, 462, 299, 321, 430, 488, 191, 404, 358, 401, 431, 517, 208, 202, 661, 153, 291, 355, 455, 555, 528, 593, 305, 433, 550, 487, 342, 574, 720, 678, 215, 314]], 'time-4': 130.64850091934204}\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wblben0Pbtwq"
      },
      "source": [
        "### LLYMask"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ye8zNyWWby_Z"
      },
      "outputs": [],
      "source": [
        "def get_model_lly(to_select):\n",
        "  model = LiaoLattyYangModel(\n",
        "      num_classes=num_classes,\n",
        "      layer_sequence=[67],\n",
        "      num_inputs=num_features,\n",
        "  )\n",
        "\n",
        "  model.compile(\n",
        "      loss=tf.keras.losses.CategoricalCrossentropy(), metrics=[\"accuracy\"]\n",
        "  )\n",
        "\n",
        "  return model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SFWaX5oHb2FW"
      },
      "outputs": [],
      "source": [
        "def get_sparse_model_evaluation_lly(mask):\n",
        "  mlp_fit = SparseModel(\n",
        "      selected_features=tf.math.reduce_sum(\n",
        "          tf.one_hot(mask, num_features, dtype=tf.float32), 0\n",
        "      ),\n",
        "      num_classes=num_classes,\n",
        "      is_classification=True,\n",
        "      layer_sequence=[67],\n",
        "  )\n",
        "  mlp_fit.compile(loss=loss_fn, metrics=[\"accuracy\"])\n",
        "  mlp_fit.fit(\n",
        "      ds_train, validation_data=ds_val, epochs=num_epochs_fit, verbose=None\n",
        "  )\n",
        "  return mlp_fit.evaluate(ds_test)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rTP0Iy0Ib387"
      },
      "outputs": [],
      "source": [
        "def get_masks_fn_lly():\n",
        "  masks = []\n",
        "  for feature_count in feature_counts:\n",
        "    print(f\"\\_ Getting mask for {feature_count} features.\")\n",
        "    mask = []\n",
        "\n",
        "    model = get_model_lly(feature_count)\n",
        "    model.fit(\n",
        "        ds_train, validation_data=ds_val, epochs=num_epochs_select, verbose=None\n",
        "    )\n",
        "\n",
        "    attention_logits = model.lly(tf.convert_to_tensor(x_train))\n",
        "    _, selected_indices = tf.math.top_k(attention_logits, k=feature_count)\n",
        "    selected_indices = selected_indices.numpy()\n",
        "\n",
        "    masks.append(list(selected_indices))\n",
        "\n",
        "  return masks"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jd2UvHlUcC8b"
      },
      "outputs": [],
      "source": [
        "# Reload global variables.\n",
        "(\n",
        "    ds_train,\n",
        "    x_train,\n",
        "    ds_val,\n",
        "    ds_test,\n",
        "    is_classification,\n",
        "    num_classes,\n",
        "    num_features,\n",
        "    num_train_steps_select,\n",
        "    loss_fn,\n",
        "    num_epochs_select,\n",
        ") = load_dataset(num_epochs_select=50, batch_size=256, dataset='mnist')\n",
        "lly_experiment_data = run_experiment(\n",
        "    get_masks_fn_lly, get_sparse_model_evaluation_lly\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v3Pm9vmkeb0V"
      },
      "outputs": [],
      "source": [
        "print(lly_experiment_data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lWEnEougmP9J"
      },
      "outputs": [],
      "source": [
        "lly_experiment_data = {0: [[0.4351731538772583, 0.8616999983787537], [0.27776700258255005, 0.9158999919891357], [0.25081080198287964, 0.9230999946594238], [0.26245900988578796, 0.9275000095367432]], 'masks-0': [[350, 400, 404, 381, 151, 542, 269, 241, 489, 523, 567, 99, 437, 629, 434, 569, 180, 237, 156, 298], [654, 268, 354, 211, 152, 486, 461, 349, 346, 456, 437, 406, 155, 516, 540, 595, 182, 568, 517, 291, 657, 518, 215, 656, 241, 272, 492, 262, 626, 242], [375, 515, 269, 322, 490, 151, 457, 711, 657, 433, 351, 183, 567, 409, 346, 541, 597, 413, 100, 415, 464, 428, 550, 377, 544, 716, 400, 191, 102, 270, 289, 290, 680, 244, 416, 378, 297, 713, 628, 573], [322, 177, 376, 525, 488, 406, 398, 595, 436, 490, 484, 712, 401, 351, 345, 438, 598, 569, 414, 155, 156, 291, 150, 152, 127, 99, 264, 153, 180, 579, 326, 550, 630, 657, 318, 297, 577, 263, 656, 578, 551, 710, 556, 247, 179, 316, 574, 304, 269, 715]], 'time-0': 110.29539012908936, 1: [[0.38512122631073, 0.8708999752998352], [0.25951626896858215, 0.9243999719619751], [0.2498072236776352, 0.9287999868392944], [0.24856728315353394, 0.9368000030517578]], 'masks-1': [[429, 399, 270, 346, 152, 517, 486, 155, 351, 597, 263, 382, 436, 405, 482, 386, 655, 430, 240, 653], [375, 524, 544, 428, 401, 461, 184, 518, 483, 318, 596, 298, 378, 181, 351, 513, 99, 124, 238, 264, 102, 408, 434, 288, 343, 243, 241, 437, 542, 433], [238, 434, 382, 514, 376, 511, 353, 464, 322, 156, 403, 296, 570, 550, 345, 516, 567, 318, 101, 97, 358, 485, 430, 177, 458, 181, 237, 460, 655, 461, 457, 182, 239, 153, 487, 435, 406, 176, 315, 684], [375, 460, 489, 322, 515, 464, 485, 323, 456, 409, 98, 183, 405, 496, 596, 318, 386, 373, 103, 150, 326, 348, 125, 382, 154, 655, 600, 466, 542, 571, 379, 269, 151, 294, 350, 207, 378, 380, 545, 247, 759, 656, 572, 211, 658, 683, 374, 626, 654, 541]], 'time-1': 105.03278303146362, 2: [[0.45899808406829834, 0.8416000008583069], [0.2820586562156677, 0.9158999919891357], [0.24951669573783875, 0.9294000267982483], [0.2521074116230011, 0.9314000010490417]], 'masks-2': [[437, 295, 427, 655, 351, 485, 571, 157, 153, 434, 405, 567, 124, 318, 429, 549, 711, 386, 517, 489], [212, 264, 322, 100, 461, 374, 155, 319, 596, 407, 289, 542, 378, 462, 512, 456, 510, 240, 211, 330, 656, 299, 270, 241, 292, 657, 381, 153, 516, 123], [486, 269, 182, 348, 462, 154, 263, 515, 570, 374, 455, 492, 429, 378, 568, 386, 381, 239, 657, 404, 655, 295, 623, 150, 538, 629, 513, 322, 579, 100, 180, 572, 72, 409, 181, 626, 464, 186, 318, 439], [270, 430, 405, 624, 542, 400, 153, 290, 350, 381, 462, 539, 379, 464, 155, 455, 489, 427, 656, 239, 267, 98, 554, 375, 294, 579, 515, 295, 177, 315, 158, 326, 740, 654, 328, 402, 628, 316, 221, 265, 513, 583, 485, 296, 709, 297, 659, 553, 240, 268]], 'time-2': 114.6936149597168, 3: [[0.38757357001304626, 0.8784999847412109], [0.305998831987381, 0.9103000164031982], [0.26736629009246826, 0.9253000020980835], [0.26541879773139954, 0.9359999895095825]], 'masks-3': [[239, 542, 183, 151, 457, 373, 263, 180, 595, 348, 569, 489, 406, 399, 378, 270, 268, 381, 298, 712], [409, 101, 544, 242, 346, 371, 487, 523, 405, 351, 434, 568, 490, 512, 154, 180, 358, 185, 570, 156, 541, 682, 436, 211, 212, 402, 400, 581, 625, 217], [297, 655, 436, 488, 154, 346, 518, 295, 399, 378, 458, 511, 541, 373, 598, 496, 150, 386, 241, 595, 375, 354, 270, 405, 381, 456, 509, 635, 244, 712, 433, 656, 243, 157, 180, 206, 203, 652, 582, 528], [376, 378, 323, 515, 241, 102, 434, 517, 467, 151, 596, 182, 442, 409, 540, 484, 319, 655, 157, 629, 326, 377, 514, 247, 401, 550, 658, 249, 464, 345, 578, 353, 713, 402, 212, 324, 743, 263, 210, 542, 101, 430, 99, 211, 582, 628, 185, 315, 682, 215]], 'time-3': 108.16814184188843, 4: [[0.385217547416687, 0.8751000165939331], [0.2757752537727356, 0.914900004863739], [0.24277658760547638, 0.9322999715805054], [0.271158903837204, 0.9291999936103821]], 'masks-4': [[297, 238, 513, 376, 290, 181, 486, 427, 381, 436, 626, 430, 101, 489, 567, 578, 657, 324, 710, 326], [183, 241, 375, 437, 326, 462, 378, 291, 540, 153, 624, 456, 488, 514, 415, 385, 538, 544, 324, 432, 373, 572, 573, 237, 556, 405, 401, 577, 124, 271], [292, 462, 377, 373, 597, 354, 512, 542, 125, 408, 456, 351, 429, 386, 155, 150, 318, 178, 657, 431, 319, 239, 577, 294, 375, 268, 482, 347, 203, 465, 538, 523, 183, 433, 176, 654, 565, 550, 148, 656], [655, 518, 403, 379, 543, 434, 512, 345, 458, 597, 515, 153, 350, 103, 538, 436, 596, 155, 386, 150, 490, 269, 186, 267, 212, 184, 214, 270, 121, 264, 491, 581, 526, 189, 147, 685, 205, 629, 353, 627, 712, 220, 100, 176, 741, 123, 159, 352, 575, 380]], 'time-4': 111.27120733261108}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-wZY12hSeeSE"
      },
      "source": [
        "### Group Lasso"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "i_qCILsXegI3"
      },
      "outputs": [],
      "source": [
        "def get_model_group_lasso(to_select):\n",
        "  model = SequentialLASSOModel(\n",
        "      num_features,\n",
        "      to_select,\n",
        "      num_train_steps=num_train_steps_select,\n",
        "      num_classes=num_classes,\n",
        "      layer_sequence=[67],\n",
        "      group_lasso_scale=0.01,\n",
        "  )\n",
        "\n",
        "  model.compile(\n",
        "      loss=tf.keras.losses.CategoricalCrossentropy(), metrics=[\"accuracy\"]\n",
        "  )\n",
        "\n",
        "  return model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aXMGKmRFegI5"
      },
      "outputs": [],
      "source": [
        "def get_sparse_model_evaluation_group_lasso(mask):\n",
        "  mlp_fit = SparseModel(\n",
        "      selected_features=tf.math.reduce_sum(\n",
        "          tf.one_hot(mask, num_features, dtype=tf.float32), 0\n",
        "      ),\n",
        "      num_classes=num_classes,\n",
        "      is_classification=True,\n",
        "      layer_sequence=[67],\n",
        "  )\n",
        "  mlp_fit.compile(loss=loss_fn, metrics=[\"accuracy\"])\n",
        "  mlp_fit.fit(\n",
        "      ds_train, validation_data=ds_val, epochs=num_epochs_fit, verbose=None\n",
        "  )\n",
        "  return mlp_fit.evaluate(ds_test)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-N466sUUegI6"
      },
      "outputs": [],
      "source": [
        "def get_masks_fn_group_lasso():\n",
        "  masks = []\n",
        "  for feature_count in feature_counts:\n",
        "    print(f\"\\_ Getting mask for {feature_count} features.\")\n",
        "    mask = []\n",
        "\n",
        "    model = get_model_group_lasso(feature_count)\n",
        "    model.fit(\n",
        "        ds_train, validation_data=ds_val, epochs=num_epochs_select, verbose=None\n",
        "    )\n",
        "\n",
        "    selected_indices = model.seql.selected_features_history.numpy().tolist()\n",
        "\n",
        "    masks.append(list(selected_indices))\n",
        "\n",
        "  return masks"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jLrIkN88egI6"
      },
      "outputs": [],
      "source": [
        "# Reload global variables.\n",
        "(\n",
        "    ds_train,\n",
        "    x_train,\n",
        "    ds_val,\n",
        "    ds_test,\n",
        "    is_classification,\n",
        "    num_classes,\n",
        "    num_features,\n",
        "    num_train_steps_select,\n",
        "    loss_fn,\n",
        "    num_epochs_select,\n",
        ") = load_dataset(num_epochs_select=50, batch_size=256, dataset='mnist')\n",
        "group_lasso_experiment_data = run_experiment(\n",
        "    get_masks_fn_group_lasso, get_sparse_model_evaluation_group_lasso\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bgaVVbnaegI7"
      },
      "outputs": [],
      "source": [
        "print(group_lasso_experiment_data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8HTmg95-mb_g"
      },
      "outputs": [],
      "source": [
        "group_lasso_experiment_data = {0: [[0.4068024456501007, 0.8687000274658203], [0.25269147753715515, 0.9232000112533569], [0.2245451956987381, 0.9388999938964844], [0.2549870014190674, 0.9488999843597412]], 'masks-0': [[350, 434, 211, 321, 514, 358, 328, 490, 379, 607, 549, 318, 348, 403, 459, 316, 516, 427, 267, 374], [350, 434, 514, 358, 321, 328, 379, 211, 342, 490, 267, 607, 316, 459, 549, 637, 516, 403, 348, 298, 323, 318, 464, 427, 405, 485, 244, 411, 269, 247], [350, 434, 358, 328, 211, 321, 490, 514, 516, 267, 379, 657, 276, 298, 607, 342, 549, 348, 459, 403, 316, 323, 318, 405, 244, 510, 427, 485, 269, 637, 572, 374, 464, 411, 263, 213, 360, 429, 570, 469], [350, 434, 358, 211, 328, 321, 514, 490, 516, 342, 348, 572, 607, 405, 459, 379, 549, 403, 298, 316, 267, 570, 427, 323, 510, 485, 318, 411, 263, 360, 464, 374, 429, 637, 247, 269, 231, 568, 244, 372, 153, 602, 469, 354, 622, 628, 657, 408, 209, 352]], 'time-0': 121.94543385505676, 1: [[0.4050743579864502, 0.869700014591217], [0.29250460863113403, 0.9053999781608582], [0.2143886685371399, 0.9442999958992004], [0.25013381242752075, 0.946399986743927]], 'masks-1': [[350, 434, 211, 321, 514, 379, 490, 607, 342, 358, 328, 316, 516, 459, 269, 318, 549, 405, 427, 572], [350, 434, 267, 358, 321, 328, 490, 379, 514, 516, 348, 211, 342, 459, 403, 323, 464, 427, 316, 269, 405, 318, 411, 247, 372, 549, 570, 637, 360, 572], [434, 350, 358, 211, 328, 321, 490, 514, 342, 379, 376, 516, 316, 298, 607, 318, 549, 267, 459, 347, 405, 323, 263, 485, 464, 411, 572, 510, 637, 554, 570, 401, 244, 360, 403, 213, 269, 427, 568, 469], [434, 350, 211, 358, 328, 490, 321, 514, 516, 267, 376, 572, 607, 405, 347, 379, 231, 298, 459, 342, 316, 570, 323, 427, 247, 637, 510, 411, 318, 360, 464, 269, 403, 485, 263, 549, 372, 244, 304, 657, 213, 429, 568, 469, 602, 354, 292, 408, 659, 622]], 'time-1': 126.89954829216003, 2: [[0.5106425881385803, 0.8226000070571899], [0.27703404426574707, 0.9133999943733215], [0.2360173612833023, 0.9366000294685364], [0.23816636204719543, 0.9509000182151794]], 'masks-2': [[350, 434, 358, 321, 328, 514, 379, 490, 342, 607, 427, 316, 298, 318, 459, 374, 348, 488, 637, 360], [434, 350, 358, 211, 328, 321, 514, 379, 490, 516, 267, 269, 348, 607, 549, 459, 316, 403, 342, 427, 323, 318, 464, 637, 411, 405, 372, 572, 374, 360], [350, 434, 358, 321, 328, 514, 490, 379, 342, 516, 459, 231, 211, 323, 607, 485, 316, 267, 347, 510, 269, 572, 318, 298, 464, 411, 403, 637, 427, 263, 360, 405, 549, 570, 244, 376, 247, 574, 401, 469], [350, 434, 358, 211, 328, 321, 514, 490, 516, 379, 348, 298, 607, 403, 459, 549, 231, 485, 267, 342, 316, 572, 360, 427, 323, 637, 318, 570, 263, 372, 464, 374, 429, 405, 510, 411, 657, 568, 244, 304, 269, 628, 467, 353, 247, 213, 292, 602, 209, 622]], 'time-2': 123.0556287765503, 3: [[0.48797357082366943, 0.8277000188827515], [0.25377538800239563, 0.9214000105857849], [0.21537768840789795, 0.9398999810218811], [0.24243970215320587, 0.9466999769210815]], 'masks-3': [[434, 350, 358, 490, 328, 321, 379, 607, 514, 298, 516, 318, 316, 427, 459, 342, 405, 348, 637, 374], [350, 434, 267, 358, 321, 328, 514, 379, 490, 211, 342, 516, 316, 459, 348, 427, 374, 269, 323, 318, 464, 637, 244, 411, 372, 549, 607, 485, 570, 572], [350, 434, 358, 328, 211, 321, 514, 379, 267, 342, 490, 657, 316, 298, 318, 405, 348, 459, 607, 516, 403, 323, 263, 485, 549, 510, 360, 637, 427, 247, 374, 572, 244, 372, 411, 213, 269, 464, 429, 570], [350, 434, 358, 211, 328, 276, 321, 348, 490, 514, 516, 379, 342, 403, 459, 607, 657, 231, 298, 549, 316, 267, 405, 318, 427, 510, 323, 485, 263, 360, 464, 374, 572, 554, 637, 269, 429, 411, 570, 372, 244, 628, 467, 568, 354, 247, 213, 408, 659, 622]], 'time-3': 136.6428382396698, 4: [[0.4293680191040039, 0.8572999835014343], [0.2715790569782257, 0.9189000129699707], [0.22757494449615479, 0.9373000264167786], [0.26412415504455566, 0.9462000131607056]], 'masks-4': [[350, 434, 211, 321, 514, 358, 328, 490, 379, 607, 459, 318, 348, 403, 427, 342, 316, 516, 269, 374], [350, 434, 321, 211, 514, 358, 328, 490, 379, 516, 348, 459, 267, 403, 342, 427, 316, 269, 464, 323, 405, 318, 411, 637, 554, 485, 549, 510, 607, 572], [350, 434, 358, 321, 328, 516, 490, 514, 267, 379, 403, 657, 211, 298, 607, 459, 348, 510, 405, 485, 427, 323, 318, 572, 464, 342, 411, 316, 637, 554, 374, 429, 372, 570, 360, 269, 628, 263, 304, 469], [350, 434, 211, 358, 328, 321, 514, 376, 379, 267, 490, 298, 405, 347, 323, 516, 657, 607, 459, 342, 549, 572, 427, 485, 316, 637, 231, 570, 244, 360, 464, 318, 269, 403, 247, 628, 510, 411, 263, 304, 401, 213, 554, 568, 602, 354, 565, 408, 659, 469]], 'time-4': 123.1732931137085}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mKpb27iQ1Kzn"
      },
      "source": [
        "## Plots (mnist)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fCnYqmr_lsXK"
      },
      "outputs": [],
      "source": [
        "experiments_data = [\n",
        "    get_data_from_experiment_output(simple_sa_experiment_data),\n",
        "    get_data_from_experiment_output(simple_sa_20_experiment_data),\n",
        "    get_data_from_experiment_output(simple_sa_40_experiment_data),\n",
        "    get_data_from_experiment_output(simple_sa_60_experiment_data),\n",
        "    get_data_from_experiment_output(\n",
        "        intermediate_sa_experiment_data, divide_time_by=5\n",
        "    ),  # ensemble of 5 models that can run in parallel\n",
        "    get_data_from_experiment_output(lly_experiment_data),\n",
        "    get_data_from_experiment_output(group_lasso_experiment_data),\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_cVcO9AghX9l"
      },
      "outputs": [],
      "source": [
        "fig = plt.gcf()\n",
        "fig.set_size_inches(18.5, 12.5)\n",
        "\n",
        "colors_for_feature_count = ['b', 'm', 'y', 'g']\n",
        "\n",
        "method_x = []\n",
        "method_y_mean = []\n",
        "method_y_std = []\n",
        "colors = []\n",
        "\n",
        "for feature_count_i in range(len(feature_counts)):\n",
        "  feature_count = feature_counts[feature_count_i]\n",
        "\n",
        "  for i in range(len(experiments_data)):\n",
        "    method = methods[i]\n",
        "    experiment_data = experiments_data[i]\n",
        "\n",
        "    method_x.append(f'{method} ({feature_count}f)')\n",
        "    method_y_mean.append(experiment_data['mean_acc_for_mask'][feature_count])\n",
        "    method_y_std.append(experiment_data['std_acc_for_mask'][feature_count])\n",
        "    colors.append(colors_for_feature_count[feature_count_i])\n",
        "\n",
        "method_x_pos = np.arange(len(method_x))\n",
        "\n",
        "fig, ax = plt.subplots()\n",
        "barlist = ax.bar(\n",
        "    method_x,\n",
        "    method_y_mean,\n",
        "    yerr=method_y_std,\n",
        "    align='center',\n",
        "    alpha=0.5,\n",
        "    ecolor='black',\n",
        "    capsize=10,\n",
        ")\n",
        "\n",
        "for i in range(len(barlist)):\n",
        "  barlist[i].set_color(colors[i])\n",
        "\n",
        "\n",
        "ax.set_xticks(method_x_pos)\n",
        "ax.set_xticklabels(method_x)\n",
        "ax.yaxis.grid(True)\n",
        "\n",
        "plt.ylabel('Accuracy')\n",
        "plt.xlabel('Method and feature set size')\n",
        "plt.xticks(rotation=30)\n",
        "ax.set_ylim([0.82, 0.96])\n",
        "plt.rcParams['figure.figsize'] = (22, 9)\n",
        "plt.title(\n",
        "    'Accuracy (mean and std) for different feature selection methods, with'\n",
        "    ' different feature set size targets, groupped by feature set size (mnist).'\n",
        ")\n",
        "\n",
        "red_patch = mpatches.Patch(color='red', label='The red data')\n",
        "blue_patch = mpatches.Patch(color='blue', label='The blue data')\n",
        "\n",
        "plt.legend(\n",
        "    handles=[\n",
        "        *[\n",
        "            mpatches.Patch(\n",
        "                color=color,\n",
        "                label=(\n",
        "                    f'{feature_count} features'\n",
        "                    f' ({round(feature_count / num_features, 4)*100}%)'\n",
        "                ),\n",
        "            )\n",
        "            for (color, feature_count) in zip(\n",
        "                colors_for_feature_count, feature_counts\n",
        "            )\n",
        "        ],\n",
        "        mpatches.Patch(\n",
        "            color='black',\n",
        "            label=(\n",
        "                'SimplaSA\u003cX\u003e: Sequential Attention on a uniform sample of X% of'\n",
        "                ' the datapoints.'\n",
        "            ),\n",
        "        ),\n",
        "    ]\n",
        ")\n",
        "\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oJD-ORI7rS1J"
      },
      "outputs": [],
      "source": [
        "fig = plt.gcf()\n",
        "fig.set_size_inches(18.5, 12.5)\n",
        "\n",
        "method_x = []\n",
        "method_y_mean = []\n",
        "method_y_std = []\n",
        "\n",
        "for i in range(len(experiments_data)):\n",
        "  method = methods[i]\n",
        "  experiment_data = experiments_data[i]\n",
        "\n",
        "  method_x.append(method)\n",
        "  method_y_mean.append(experiment_data['mean_time'])\n",
        "  method_y_std.append(experiment_data['std_time'])\n",
        "\n",
        "method_x_pos = np.arange(len(method_x))\n",
        "\n",
        "fig, ax = plt.subplots()\n",
        "barlist = ax.bar(\n",
        "    method_x,\n",
        "    method_y_mean,\n",
        "    yerr=method_y_std,\n",
        "    align='center',\n",
        "    alpha=0.5,\n",
        "    ecolor='black',\n",
        "    capsize=10,\n",
        ")\n",
        "\n",
        "ax.set_xticks(method_x_pos)\n",
        "ax.set_xticklabels(method_x)\n",
        "ax.yaxis.grid(True)\n",
        "\n",
        "plt.ylabel('Seconds')\n",
        "plt.xlabel('Method')\n",
        "plt.rcParams['figure.figsize'] = (22, 9)\n",
        "plt.title('End-to-end time to compute masks for all different sizes (mnist).')\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rCZMwu9rxUqM"
      },
      "outputs": [],
      "source": [
        "divide_by = experiments_data[0]['mean_time']\n",
        "\n",
        "fig = plt.gcf()\n",
        "fig.set_size_inches(18.5, 12.5)\n",
        "\n",
        "method_x = []\n",
        "method_y_mean = []\n",
        "method_y_std = []\n",
        "\n",
        "for i in range(len(experiments_data)):\n",
        "  method = methods[i]\n",
        "  experiment_data = experiments_data[i]\n",
        "\n",
        "  method_x.append(method)\n",
        "  method_y_mean.append(experiment_data['mean_time'] / divide_by)\n",
        "  method_y_std.append(experiment_data['std_time'] / divide_by)\n",
        "\n",
        "method_x_pos = np.arange(len(method_x))\n",
        "\n",
        "fig, ax = plt.subplots()\n",
        "barlist = ax.bar(\n",
        "    method_x,\n",
        "    method_y_mean,\n",
        "    yerr=method_y_std,\n",
        "    align='center',\n",
        "    alpha=0.5,\n",
        "    ecolor='black',\n",
        "    capsize=10,\n",
        ")\n",
        "\n",
        "ax.set_xticks(method_x_pos)\n",
        "ax.set_xticklabels(method_x)\n",
        "ax.yaxis.grid(True)\n",
        "\n",
        "plt.ylabel('Time relative to simple Sequential Attention')\n",
        "plt.xlabel('Method')\n",
        "plt.rcParams['figure.figsize'] = (22, 9)\n",
        "plt.title('Time to compute masks for all different sizes (mnist).')\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eURAvRz92kwV"
      },
      "outputs": [],
      "source": [
        "divide_by = (\n",
        "    experiment_data['mean_acc_for_mask'][feature_counts[-1]]\n",
        "    / experiment_data['mean_time']\n",
        ")\n",
        "fig = plt.gcf()\n",
        "fig.set_size_inches(18.5, 12.5)\n",
        "\n",
        "method_x = []\n",
        "method_y_mean = []\n",
        "method_y_std = []\n",
        "\n",
        "for i in range(len(experiments_data)):\n",
        "  method = methods[i]\n",
        "  experiment_data = experiments_data[i]\n",
        "\n",
        "  method_x.append(method)\n",
        "  method_y_mean.append(\n",
        "      experiment_data['mean_acc_for_mask'][feature_counts[-1]]\n",
        "      / experiment_data['mean_time']\n",
        "      / divide_by\n",
        "  )\n",
        "  method_y_std.append(\n",
        "      experiment_data['std_acc_for_mask'][feature_counts[-1]] / divide_by\n",
        "  )\n",
        "\n",
        "method_x_pos = np.arange(len(method_x))\n",
        "\n",
        "fig, ax = plt.subplots()\n",
        "barlist = ax.bar(\n",
        "    method_x,\n",
        "    method_y_mean,\n",
        "    yerr=method_y_std,\n",
        "    align='center',\n",
        "    alpha=0.5,\n",
        "    ecolor='black',\n",
        "    capsize=10,\n",
        ")\n",
        "\n",
        "ax.set_xticks(method_x_pos)\n",
        "ax.set_xticklabels(method_x)\n",
        "ax.yaxis.grid(True)\n",
        "\n",
        "plt.ylabel('Ratio relative to Sequential attention.')\n",
        "plt.xlabel('Method')\n",
        "plt.rcParams['figure.figsize'] = (22, 9)\n",
        "plt.title('Final accuracy / Time to find the mask ratio (50 features, mnist)')\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IEZeaevV8yyK"
      },
      "source": [
        "## Plots (mnist) - measuring the impact of random initialization on simple SA vs intermediate SA"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VzEmjxOK9DaU"
      },
      "outputs": [],
      "source": [
        "mean_acc_simple_sa = []\n",
        "std_acc_simple_sa = []\n",
        "mean_acc_intermediate_sa = []\n",
        "std_acc_intermediate_sa = []\n",
        "\n",
        "for feature_count_i in range(len(feature_counts)):\n",
        "  mean_acc_simple_sa.append(\n",
        "      np.mean(\n",
        "          [\n",
        "              simple_sa_experiment_data[key][feature_count_i][1]\n",
        "              for key in range(random_runs_count)\n",
        "          ]\n",
        "      )\n",
        "  )\n",
        "  std_acc_simple_sa.append(\n",
        "      np.std(\n",
        "          [\n",
        "              simple_sa_experiment_data[key][feature_count_i][1]\n",
        "              for key in range(random_runs_count)\n",
        "          ]\n",
        "      )\n",
        "  )\n",
        "\n",
        "  mean_acc_intermediate_sa.append(\n",
        "      np.mean(\n",
        "          [\n",
        "              intermediate_sa_experiment_data[key][feature_count_i][1]\n",
        "              for key in range(random_runs_count)\n",
        "          ]\n",
        "      )\n",
        "  )\n",
        "  std_acc_intermediate_sa.append(\n",
        "      np.std(\n",
        "          [\n",
        "              intermediate_sa_experiment_data[key][feature_count_i][1]\n",
        "              for key in range(random_runs_count)\n",
        "          ]\n",
        "      )\n",
        "  )\n",
        "\n",
        "mean_acc_simple_sa = np.array(mean_acc_simple_sa)\n",
        "std_acc_simple_sa = np.array(std_acc_simple_sa)\n",
        "mean_acc_intermediate_sa = np.array(mean_acc_intermediate_sa)\n",
        "std_acc_intermediate_sa = np.array(std_acc_intermediate_sa)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4M9Kerev-htX"
      },
      "outputs": [],
      "source": [
        "x = feature_counts\n",
        "\n",
        "plt.figure(figsize=(16, 5))\n",
        "\n",
        "left = plt.subplot(1,2,1)\n",
        "right = plt.subplot(1,2,2)\n",
        "\n",
        "left.plot(x, mean_acc_simple_sa, 'b', label='Sequential Attention')\n",
        "left.fill_between(\n",
        "    x,\n",
        "    mean_acc_simple_sa + std_acc_simple_sa,\n",
        "    mean_acc_simple_sa - std_acc_simple_sa,\n",
        "    color='#a5baf97f',\n",
        ")\n",
        "\n",
        "right.plot(\n",
        "    x, mean_acc_intermediate_sa, 'g', label='Intermediate Sequential Attention'\n",
        ")\n",
        "right.fill_between(\n",
        "    x,\n",
        "    mean_acc_intermediate_sa + std_acc_intermediate_sa,\n",
        "    mean_acc_intermediate_sa - std_acc_intermediate_sa,\n",
        "    color='#c1f4ba7f',\n",
        ")\n",
        "\n",
        "left.title.set_text('Simple Sequential Attention (mnist)')\n",
        "right.title.set_text('Intermediate Sequential Attention (mnist)')\n",
        "right.set_xlabel('Feature count')\n",
        "right.set_ylabel('Accuracy')\n",
        "left.set_xlabel('Feature count')\n",
        "left.set_ylabel('Accuracy')\n",
        "\n",
        "ax = plt.gca()\n",
        "ax.set_axisbelow(True)\n",
        "\n",
        "left.grid(True)\n",
        "right.grid(True)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ulra3oKXGNoF"
      },
      "source": [
        "## Time comparisons for different data quantities (mnist)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "43DLefOkGvsI"
      },
      "source": [
        "### 100% - simple SA\n",
        "Sequential attention on 100% of the dataset (randomly sampled)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o266_aCHGxRy"
      },
      "outputs": [],
      "source": [
        "num_epochs_select = 50\n",
        "datasets = get_dataset(0.125, 256, subsample_size=1.0)\n",
        "x_train = datasets[\"x_train\"]\n",
        "ds_train = datasets[\"ds_train\"]\n",
        "ds_val = datasets[\"ds_val\"]\n",
        "ds_test = datasets[\"ds_test\"]\n",
        "is_classification = datasets[\"is_classification\"]\n",
        "num_classes = datasets[\"num_classes\"]\n",
        "num_features = datasets[\"num_features\"]\n",
        "num_train_steps_select = num_epochs_select * len(ds_train)\n",
        "loss_fn = (\n",
        "    tf.keras.losses.CategoricalCrossentropy()\n",
        "    if is_classification\n",
        "    else tf.keras.losses.MeanAbsoluteError()\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GNmbWtS0G-sE"
      },
      "outputs": [],
      "source": [
        "model = get_model_simple_sa(50)\n",
        "start = time.time()\n",
        "\n",
        "model.fit(ds_train, validation_data=ds_val, epochs=num_epochs_select, verbose=2)\n",
        "\n",
        "end = time.time()\n",
        "print(f\"Time 100%: {end - start}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1X1Vy6-8HXma"
      },
      "source": [
        "### 20% - simple SA\n",
        "Sequential attention on 20% of the dataset (randomly sampled)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fxVJCH0JbebM"
      },
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1t7fLoDHHZ4J"
      },
      "outputs": [],
      "source": [
        "num_epochs_select = 50\n",
        "datasets = get_dataset(0.125, 256, subsample_size=0.2)\n",
        "x_train = datasets[\"x_train\"]\n",
        "ds_train = datasets[\"ds_train\"]\n",
        "ds_val = datasets[\"ds_val\"]\n",
        "ds_test = datasets[\"ds_test\"]\n",
        "is_classification = datasets[\"is_classification\"]\n",
        "num_classes = datasets[\"num_classes\"]\n",
        "num_features = datasets[\"num_features\"]\n",
        "num_train_steps_select = num_epochs_select * len(ds_train)\n",
        "loss_fn = (\n",
        "    tf.keras.losses.CategoricalCrossentropy()\n",
        "    if is_classification\n",
        "    else tf.keras.losses.MeanAbsoluteError()\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zXiSHia3HdD8"
      },
      "outputs": [],
      "source": [
        "model = get_model_simple_sa(50)\n",
        "start = time.time()\n",
        "\n",
        "model.fit(ds_train, validation_data=ds_val, epochs=num_epochs_select, verbose=2)\n",
        "\n",
        "end = time.time()\n",
        "print(f\"Time 20%: {end - start}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7WUIfDTLc6Bc"
      },
      "source": [
        "### 40% - simple SA\n",
        "Sequential attention on 40% of the dataset (randomly sampled)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uCczsRuwc3-Y"
      },
      "outputs": [],
      "source": [
        "num_epochs_select = 50\n",
        "datasets = get_dataset(0.125, 256, subsample_size=0.4)\n",
        "x_train = datasets[\"x_train\"]\n",
        "ds_train = datasets[\"ds_train\"]\n",
        "ds_val = datasets[\"ds_val\"]\n",
        "ds_test = datasets[\"ds_test\"]\n",
        "is_classification = datasets[\"is_classification\"]\n",
        "num_classes = datasets[\"num_classes\"]\n",
        "num_features = datasets[\"num_features\"]\n",
        "num_train_steps_select = num_epochs_select * len(ds_train)\n",
        "loss_fn = (\n",
        "    tf.keras.losses.CategoricalCrossentropy()\n",
        "    if is_classification\n",
        "    else tf.keras.losses.MeanAbsoluteError()\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2H-Bd5oXc3-a"
      },
      "outputs": [],
      "source": [
        "model = get_model_simple_sa(50)\n",
        "start = time.time()\n",
        "\n",
        "model.fit(ds_train, validation_data=ds_val, epochs=num_epochs_select, verbose=2)\n",
        "\n",
        "end = time.time()\n",
        "print(f\"Time 40%: {end - start}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8y9IWCOMdGBL"
      },
      "source": [
        "### 60% - simple SA\n",
        "Sequential attention on 60% of the dataset (randomly sampled)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WO9yt8DxdGBN"
      },
      "outputs": [],
      "source": [
        "num_epochs_select = 50\n",
        "datasets = get_dataset(0.125, 256, subsample_size=0.6)\n",
        "x_train = datasets[\"x_train\"]\n",
        "ds_train = datasets[\"ds_train\"]\n",
        "ds_val = datasets[\"ds_val\"]\n",
        "ds_test = datasets[\"ds_test\"]\n",
        "is_classification = datasets[\"is_classification\"]\n",
        "num_classes = datasets[\"num_classes\"]\n",
        "num_features = datasets[\"num_features\"]\n",
        "num_train_steps_select = num_epochs_select * len(ds_train)\n",
        "loss_fn = (\n",
        "    tf.keras.losses.CategoricalCrossentropy()\n",
        "    if is_classification\n",
        "    else tf.keras.losses.MeanAbsoluteError()\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XFFakfmUdGBP"
      },
      "outputs": [],
      "source": [
        "model = get_model_simple_sa(50)\n",
        "start = time.time()\n",
        "\n",
        "model.fit(ds_train, validation_data=ds_val, epochs=num_epochs_select, verbose=2)\n",
        "\n",
        "end = time.time()\n",
        "print(f\"Time 60%: {end - start}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rivHjiQv587w"
      },
      "source": [
        "## Timed runs (Activity)\n",
        "Goal: track the time it takes to get a mask, then track the performance of a sparse model trained with it"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GdonklQQ2XyB"
      },
      "outputs": [],
      "source": [
        "# Mnist config.\n",
        "num_epochs_fit = 134"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BsvB8nFm587y"
      },
      "source": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PV4HRilI587z"
      },
      "source": [
        "### Simple SA"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Bq3CwfVO587z"
      },
      "outputs": [],
      "source": [
        "def get_model_simple_sa(to_select):\n",
        "  model = SequentialAttentionMlp(\n",
        "      num_train_steps=num_train_steps_select,\n",
        "      num_candidates_to_select=to_select,\n",
        "      num_classes=num_classes,\n",
        "      layer_sequence=[67],\n",
        "      learning_rate=1e-5,\n",
        "      decay_steps=500,\n",
        "      decay_rate=1.0,\n",
        "      num_candidates=num_features,\n",
        "  )\n",
        "\n",
        "  model.compile(\n",
        "      loss=tf.keras.losses.CategoricalCrossentropy(), metrics=[\"accuracy\"]\n",
        "  )\n",
        "\n",
        "  return model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o0uMlDpU587z"
      },
      "outputs": [],
      "source": [
        "def get_sparse_model_evaluation_simple_sa(mask):\n",
        "  mlp_fit = SparseModel(\n",
        "      selected_features=tf.math.reduce_sum(\n",
        "          tf.one_hot(mask, num_features, dtype=tf.float32), 0\n",
        "      ),\n",
        "      num_classes=num_classes,\n",
        "      is_classification=True,\n",
        "      layer_sequence=[67],\n",
        "  )\n",
        "  mlp_fit.compile(loss=loss_fn, metrics=[\"accuracy\"])\n",
        "  mlp_fit.fit(\n",
        "      ds_train, validation_data=ds_val, epochs=num_epochs_fit, verbose=None\n",
        "  )\n",
        "  return mlp_fit.evaluate(ds_test)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IVn7e2IX587z"
      },
      "outputs": [],
      "source": [
        "def get_masks_fn_simple_sa():\n",
        "  masks = []\n",
        "  for feature_count in feature_counts:\n",
        "    print(f\"\\_ Getting mask for {feature_count} features.\")\n",
        "    mask = []\n",
        "\n",
        "    model = get_model_simple_sa(feature_count)\n",
        "    model.fit(\n",
        "        ds_train, validation_data=ds_val, epochs=num_epochs_select, verbose=None\n",
        "    )\n",
        "\n",
        "    selected_features = model.sa.selected_features\n",
        "    _, selected_indices = tf.math.top_k(selected_features, k=feature_count)\n",
        "    selected_indices = selected_indices.numpy()\n",
        "\n",
        "    masks.append(list(selected_indices))\n",
        "\n",
        "  return masks"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t4qwkZYu587z"
      },
      "outputs": [],
      "source": [
        "# Reload global variables.\n",
        "(\n",
        "    ds_train,\n",
        "    x_train,\n",
        "    ds_val,\n",
        "    ds_test,\n",
        "    is_classification,\n",
        "    num_classes,\n",
        "    num_features,\n",
        "    num_train_steps_select,\n",
        "    loss_fn,\n",
        "    num_epochs_select,\n",
        ") = load_dataset(num_epochs_select=624, batch_size=256, dataset='activity')\n",
        "simple_sa_experiment_data = run_experiment(\n",
        "    get_masks_fn_simple_sa, get_sparse_model_evaluation_simple_sa\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y6KGPeBA587z"
      },
      "outputs": [],
      "source": [
        "print(simple_sa_experiment_data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8KDFShg-ca0b"
      },
      "outputs": [],
      "source": [
        "simple_sa_experiment_data = {0: [[0.3504304885864258, 0.9015948176383972], [0.3777216970920563, 0.9175432920455933], [0.4533860385417938, 0.9250084757804871], [0.42957890033721924, 0.935188353061676]], 'masks-0': [[37, 40, 50, 62, 63, 69, 89, 149, 159, 182, 185, 198, 222, 298, 302, 369, 451, 458, 527, 537], [9, 28, 37, 39, 40, 50, 62, 63, 69, 77, 100, 142, 149, 159, 182, 185, 186, 192, 198, 297, 298, 371, 372, 374, 450, 451, 456, 458, 508, 526], [36, 37, 39, 40, 49, 50, 53, 62, 63, 65, 70, 77, 89, 112, 117, 142, 149, 158, 159, 182, 185, 186, 187, 192, 196, 197, 198, 298, 302, 341, 369, 371, 372, 374, 450, 451, 456, 458, 461, 526], [28, 32, 36, 37, 39, 41, 52, 55, 57, 62, 63, 65, 69, 77, 79, 89, 91, 106, 112, 117, 118, 142, 149, 158, 159, 182, 185, 186, 187, 192, 196, 197, 198, 202, 222, 295, 298, 369, 371, 374, 378, 433, 450, 451, 456, 458, 461, 526, 538, 553]], 'time-0': 238.6969313621521, 1: [[0.30231568217277527, 0.9107567071914673], [0.45284411311149597, 0.8995589017868042], [0.47835686802864075, 0.9233118295669556], [0.4203539788722992, 0.9287410974502563]], 'masks-1': [[37, 39, 40, 50, 65, 69, 89, 142, 149, 159, 185, 198, 251, 295, 366, 369, 450, 451, 458, 508], [25, 36, 40, 41, 65, 72, 77, 89, 91, 112, 117, 118, 142, 149, 159, 182, 185, 186, 192, 197, 198, 298, 302, 369, 371, 374, 450, 451, 459, 526], [28, 36, 37, 39, 40, 53, 55, 62, 63, 65, 71, 77, 112, 117, 118, 142, 149, 158, 159, 185, 186, 192, 196, 197, 198, 199, 295, 299, 300, 369, 371, 424, 450, 451, 452, 456, 459, 508, 527, 553], [28, 32, 36, 37, 39, 40, 41, 50, 55, 57, 58, 62, 63, 65, 69, 77, 112, 117, 118, 119, 142, 145, 149, 157, 159, 182, 185, 186, 187, 192, 196, 197, 198, 295, 298, 302, 369, 371, 373, 374, 378, 390, 445, 451, 457, 458, 461, 508, 527, 553]], 'time-1': 254.16656279563904, 2: [[0.41085848212242126, 0.9022734761238098], [0.4047444760799408, 0.9131320118904114], [0.3805268108844757, 0.9219545125961304], [0.4773891270160675, 0.913810670375824]], 'masks-2': [[39, 40, 53, 66, 89, 117, 142, 149, 159, 185, 222, 302, 369, 371, 377, 450, 451, 457, 458, 526], [28, 36, 39, 40, 50, 62, 65, 69, 112, 117, 142, 149, 159, 182, 185, 186, 187, 192, 198, 295, 298, 369, 371, 395, 450, 451, 457, 459, 508, 526], [21, 28, 32, 36, 39, 41, 49, 62, 63, 65, 69, 77, 104, 105, 112, 114, 117, 118, 142, 149, 159, 182, 185, 186, 192, 196, 197, 198, 369, 374, 376, 377, 450, 451, 457, 459, 461, 508, 527, 553], [28, 32, 36, 37, 39, 49, 50, 57, 62, 63, 65, 69, 70, 77, 100, 112, 117, 118, 119, 142, 149, 159, 176, 182, 185, 186, 187, 192, 196, 197, 198, 222, 298, 302, 310, 360, 369, 371, 374, 378, 397, 450, 451, 456, 458, 461, 489, 508, 526, 553]], 'time-2': 253.64411640167236, 3: [[0.30259114503860474, 0.9110960364341736], [0.33208131790161133, 0.9277231097221375], [0.34958937764167786, 0.9300984144210815], [0.4965287148952484, 0.9209365248680115]], 'masks-3': [[37, 40, 50, 59, 69, 89, 149, 159, 182, 185, 198, 295, 298, 302, 369, 371, 450, 451, 458, 526], [22, 36, 37, 39, 49, 50, 62, 63, 69, 77, 89, 112, 142, 149, 159, 182, 185, 186, 196, 295, 298, 302, 369, 374, 450, 451, 456, 491, 526, 553], [22, 28, 32, 36, 37, 39, 49, 53, 62, 63, 65, 69, 89, 91, 105, 112, 117, 142, 145, 148, 149, 159, 182, 185, 186, 192, 196, 197, 198, 295, 298, 314, 371, 374, 450, 451, 456, 459, 527, 558], [21, 22, 32, 36, 39, 41, 45, 50, 52, 55, 62, 63, 65, 69, 77, 89, 91, 112, 117, 118, 142, 145, 149, 159, 176, 182, 185, 186, 187, 192, 196, 197, 198, 236, 295, 297, 298, 300, 314, 369, 371, 450, 451, 457, 458, 459, 461, 508, 526, 553]], 'time-3': 257.2544434070587, 4: [[0.3400125503540039, 0.9066847562789917], [0.375901460647583, 0.9093993902206421], [0.6113688349723816, 0.907024085521698], [0.4592318832874298, 0.9287410974502563]], 'masks-4': [[37, 40, 50, 62, 69, 89, 149, 159, 182, 185, 191, 198, 298, 302, 369, 374, 450, 451, 456, 459], [22, 28, 32, 37, 39, 40, 50, 62, 69, 89, 91, 112, 117, 142, 149, 159, 182, 185, 186, 192, 196, 197, 198, 302, 371, 372, 374, 450, 451, 459], [28, 32, 36, 39, 41, 49, 55, 57, 62, 71, 77, 89, 112, 117, 142, 149, 158, 159, 185, 186, 187, 192, 196, 197, 198, 199, 210, 295, 298, 300, 302, 369, 371, 442, 450, 451, 456, 458, 526, 537], [28, 36, 37, 38, 39, 40, 41, 55, 57, 62, 63, 70, 77, 112, 117, 118, 142, 145, 149, 154, 158, 159, 182, 185, 186, 187, 192, 196, 197, 198, 236, 298, 302, 316, 369, 371, 372, 374, 379, 411, 433, 445, 450, 451, 456, 459, 461, 508, 526, 540]], 'time-4': 252.2806839942932}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4yUU-5KF5870"
      },
      "source": [
        "### Simple SA 20\n",
        "Sequential attention on 20% of the dataset (randomly sampled)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X4Gb1oZr5870"
      },
      "outputs": [],
      "source": [
        "def get_model_simple_sa20(to_select):\n",
        "  return get_model_simple_sa(to_select)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6yRQgO-I5870"
      },
      "outputs": [],
      "source": [
        "def get_sparse_model_evaluation_simple_sa20(mask):\n",
        "  return get_sparse_model_evaluation_simple_sa(mask)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vI6VTWOL5871"
      },
      "outputs": [],
      "source": [
        "def get_masks_fn_simple_sa20():\n",
        "  return get_masks_fn_simple_sa()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c4WyI0SN5871"
      },
      "outputs": [],
      "source": [
        "# Reload global variables.\n",
        "(\n",
        "    ds_train,\n",
        "    x_train,\n",
        "    ds_val,\n",
        "    ds_test,\n",
        "    is_classification,\n",
        "    num_classes,\n",
        "    num_features,\n",
        "    num_train_steps_select,\n",
        "    loss_fn,\n",
        "    num_epochs_select,\n",
        ") = load_dataset(\n",
        "    num_epochs_select=624,\n",
        "    batch_size=256,\n",
        "    dataset='activity',\n",
        "    subsample_size=0.2,\n",
        ")\n",
        "simple_sa_20_experiment_data = run_experiment(\n",
        "    get_masks_fn_simple_sa20, get_sparse_model_evaluation_simple_sa20\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kZ-jAMwl5871"
      },
      "outputs": [],
      "source": [
        "print(simple_sa_20_experiment_data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_UdGanoS5871"
      },
      "outputs": [],
      "source": [
        "simple_sa_20_experiment_data = {0: [[0.3777690827846527, 0.8825924396514893], [0.4948747456073761, 0.8754665851593018], [0.35029906034469604, 0.9134713411331177], [0.4793199598789215, 0.8998982310295105]], 'masks-0': [[15, 36, 40, 41, 63, 65, 77, 78, 117, 149, 185, 198, 295, 326, 377, 451, 457, 459, 508, 552], [28, 50, 52, 62, 63, 69, 77, 78, 79, 115, 117, 118, 119, 149, 185, 186, 187, 192, 197, 198, 273, 291, 310, 451, 456, 459, 461, 527, 540, 552], [28, 36, 38, 42, 46, 52, 53, 62, 63, 69, 77, 78, 103, 115, 117, 118, 119, 149, 157, 185, 186, 187, 192, 197, 198, 199, 290, 291, 295, 302, 450, 451, 456, 458, 461, 508, 526, 537, 540, 553], [28, 32, 37, 38, 41, 52, 55, 56, 62, 63, 70, 77, 78, 95, 115, 117, 118, 119, 142, 149, 157, 158, 159, 185, 186, 187, 189, 192, 193, 195, 197, 198, 199, 243, 270, 291, 314, 369, 396, 449, 450, 451, 457, 459, 482, 513, 526, 537, 552, 559]], 'time-0': 105.65340566635132, 1: [[0.31028687953948975, 0.8890396952629089], [0.5120583772659302, 0.876823902130127], [0.46726465225219727, 0.8937903046607971], [0.35787490010261536, 0.9141499996185303]], 'masks-1': [[9, 28, 61, 63, 65, 71, 77, 78, 117, 142, 149, 185, 198, 451, 457, 458, 508, 526, 552, 558], [28, 38, 49, 53, 63, 65, 69, 77, 78, 115, 117, 118, 149, 152, 185, 186, 187, 198, 199, 291, 295, 324, 450, 451, 457, 458, 508, 527, 539, 555], [13, 28, 38, 53, 62, 63, 69, 77, 115, 117, 118, 119, 142, 149, 156, 157, 159, 185, 186, 187, 197, 198, 199, 209, 291, 294, 302, 316, 320, 377, 416, 422, 449, 451, 457, 458, 527, 537, 552, 558], [3, 19, 28, 36, 38, 39, 40, 41, 42, 62, 63, 65, 69, 77, 78, 87, 117, 118, 119, 145, 149, 157, 159, 185, 186, 187, 192, 197, 198, 199, 233, 236, 249, 263, 291, 302, 315, 372, 374, 428, 449, 451, 456, 458, 461, 527, 539, 541, 553, 555]], 'time-1': 114.37213444709778, 2: [[0.417705774307251, 0.8588395118713379], [0.4132900536060333, 0.8920936584472656], [0.36045241355895996, 0.9121140241622925], [0.3518473207950592, 0.9080420732498169]], 'masks-2': [[36, 52, 62, 63, 69, 77, 78, 117, 142, 152, 183, 185, 198, 302, 374, 451, 452, 457, 459, 553], [28, 41, 56, 63, 65, 71, 78, 115, 117, 118, 142, 149, 159, 185, 186, 187, 198, 236, 278, 291, 302, 317, 451, 457, 458, 507, 527, 552, 554, 556], [1, 28, 36, 37, 39, 40, 42, 63, 65, 79, 117, 118, 157, 159, 176, 185, 186, 187, 189, 197, 198, 210, 226, 290, 291, 294, 299, 302, 374, 378, 391, 450, 451, 457, 459, 526, 537, 539, 553, 559], [26, 28, 36, 38, 50, 52, 54, 62, 63, 69, 77, 78, 117, 118, 119, 122, 126, 145, 149, 157, 159, 182, 185, 186, 187, 194, 197, 198, 236, 263, 291, 296, 308, 310, 320, 323, 374, 391, 417, 444, 450, 451, 456, 459, 463, 465, 527, 537, 540, 552]], 'time-2': 112.7014811038971, 3: [[0.346096932888031, 0.8927723169326782], [0.4475404620170593, 0.8954869508743286], [0.3728483021259308, 0.9029521346092224], [0.36161908507347107, 0.9015948176383972]], 'masks-3': [[37, 40, 53, 62, 63, 70, 78, 115, 149, 185, 197, 198, 270, 291, 308, 451, 457, 459, 461, 508], [28, 38, 39, 41, 49, 62, 63, 69, 77, 78, 111, 115, 117, 118, 142, 149, 159, 185, 186, 187, 192, 302, 374, 378, 450, 451, 457, 459, 516, 537], [28, 36, 38, 41, 47, 54, 63, 68, 70, 77, 80, 117, 118, 149, 157, 159, 185, 186, 187, 192, 198, 236, 291, 310, 316, 323, 367, 374, 377, 396, 450, 451, 456, 458, 465, 508, 537, 540, 553, 558], [28, 32, 36, 40, 43, 51, 53, 55, 62, 63, 70, 77, 78, 89, 117, 118, 119, 142, 145, 158, 159, 182, 186, 187, 189, 192, 197, 198, 199, 236, 260, 291, 294, 296, 310, 317, 374, 420, 450, 451, 456, 459, 469, 473, 488, 508, 527, 540, 552, 555]], 'time-3': 106.86495757102966, 4: [[0.3862552046775818, 0.885646402835846], [0.48700064420700073, 0.873430609703064], [0.48729416728019714, 0.8978622555732727], [0.44500991702079773, 0.8995589017868042]], 'masks-4': [[28, 36, 41, 52, 63, 69, 78, 117, 149, 185, 198, 199, 236, 274, 451, 457, 459, 461, 508, 553], [28, 38, 41, 52, 62, 63, 65, 77, 78, 115, 117, 118, 185, 187, 189, 198, 199, 232, 278, 291, 302, 316, 451, 457, 459, 461, 508, 527, 539, 552], [27, 28, 38, 50, 52, 65, 70, 78, 79, 115, 117, 118, 122, 142, 152, 159, 185, 186, 187, 189, 190, 192, 194, 197, 198, 271, 291, 295, 370, 381, 387, 450, 451, 457, 458, 465, 526, 537, 539, 552], [28, 38, 40, 42, 53, 54, 55, 62, 63, 70, 77, 78, 79, 115, 117, 118, 122, 142, 145, 149, 151, 159, 185, 186, 187, 197, 198, 236, 263, 266, 278, 291, 297, 302, 365, 370, 371, 378, 424, 444, 450, 451, 457, 459, 488, 526, 537, 540, 552, 555]], 'time-4': 111.55968928337097}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QoMom8gd5872"
      },
      "source": [
        "### Simple SA 40\n",
        "Sequential attention on 40% of the dataset (randomly sampled)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cW9yC_2t5872"
      },
      "outputs": [],
      "source": [
        "def get_model_simple_sa40(to_select):\n",
        "  return get_model_simple_sa(to_select)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vAhD4HVO5872"
      },
      "outputs": [],
      "source": [
        "def get_sparse_model_evaluation_simple_sa40(mask):\n",
        "  return get_sparse_model_evaluation_simple_sa(mask)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BA1abvWy5872"
      },
      "outputs": [],
      "source": [
        "def get_masks_fn_simple_sa40():\n",
        "  return get_masks_fn_simple_sa()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XlKQI2MF5872"
      },
      "outputs": [],
      "source": [
        "# Reload global variables.\n",
        "(\n",
        "    ds_train,\n",
        "    x_train,\n",
        "    ds_val,\n",
        "    ds_test,\n",
        "    is_classification,\n",
        "    num_classes,\n",
        "    num_features,\n",
        "    num_train_steps_select,\n",
        "    loss_fn,\n",
        "    num_epochs_select,\n",
        ") = load_dataset(\n",
        "    num_epochs_select=624,\n",
        "    batch_size=256,\n",
        "    dataset='activity',\n",
        "    subsample_size=0.4,\n",
        ")\n",
        "simple_sa_40_experiment_data = run_experiment(\n",
        "    get_masks_fn_simple_sa40, get_sparse_model_evaluation_simple_sa40\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iY5orD855873"
      },
      "outputs": [],
      "source": [
        "print(simple_sa_40_experiment_data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7ca4gP0E5873"
      },
      "outputs": [],
      "source": [
        "simple_sa_40_experiment_data = {0: [[0.33771032094955444, 0.9029521346092224], [0.4751144051551819, 0.8961656093597412], [0.4267192780971527, 0.913810670375824], [0.6220530867576599, 0.9060060977935791]], 'masks-0': [[19, 32, 41, 52, 63, 65, 77, 117, 118, 149, 159, 185, 198, 295, 377, 450, 451, 457, 459, 508], [26, 32, 36, 52, 53, 62, 63, 77, 117, 118, 149, 159, 185, 192, 197, 198, 209, 298, 302, 306, 369, 371, 378, 451, 457, 459, 508, 526, 537, 540], [28, 32, 36, 41, 52, 62, 63, 65, 69, 77, 114, 117, 118, 142, 149, 158, 185, 186, 187, 192, 196, 197, 198, 199, 222, 224, 295, 299, 300, 350, 369, 371, 396, 450, 451, 459, 489, 508, 526, 540], [28, 32, 36, 37, 49, 50, 53, 58, 62, 63, 64, 65, 70, 77, 89, 106, 112, 114, 117, 118, 119, 142, 149, 152, 158, 159, 185, 186, 187, 192, 196, 197, 198, 224, 237, 275, 287, 295, 299, 302, 371, 450, 451, 452, 457, 459, 527, 537, 540, 547]], 'time-0': 145.65792393684387, 1: [[0.2739471197128296, 0.917203962802887], [0.4335939884185791, 0.9015948176383972], [0.5463259220123291, 0.8893790245056152], [0.4611562192440033, 0.9049881100654602]], 'masks-1': [[37, 53, 62, 63, 70, 77, 149, 159, 198, 260, 295, 298, 302, 451, 457, 459, 489, 526, 537, 558], [4, 28, 32, 40, 41, 62, 63, 69, 77, 117, 118, 142, 149, 159, 185, 186, 187, 192, 196, 198, 212, 298, 302, 371, 451, 457, 459, 527, 537, 540], [16, 28, 32, 36, 49, 53, 55, 62, 63, 69, 77, 89, 117, 118, 119, 142, 149, 158, 159, 185, 186, 187, 192, 196, 197, 198, 224, 287, 295, 298, 300, 369, 371, 450, 451, 459, 461, 527, 537, 540], [28, 32, 36, 37, 40, 41, 50, 62, 63, 65, 74, 77, 89, 102, 114, 115, 117, 118, 119, 142, 145, 149, 152, 158, 159, 185, 186, 187, 192, 196, 197, 198, 222, 225, 238, 298, 310, 320, 371, 378, 386, 450, 451, 459, 461, 508, 527, 537, 540, 558]], 'time-1': 152.0008466243744, 2: [[0.3597760498523712, 0.9022734761238098], [0.43831053376197815, 0.9015948176383972], [0.46501094102859497, 0.9015948176383972], [0.39960044622421265, 0.9199185371398926]], 'masks-2': [[39, 40, 41, 62, 63, 117, 149, 159, 185, 295, 298, 302, 321, 371, 377, 451, 456, 459, 526, 537], [36, 49, 50, 62, 63, 69, 77, 114, 117, 118, 119, 142, 149, 154, 159, 185, 192, 197, 198, 211, 295, 298, 302, 374, 450, 451, 456, 459, 526, 537], [21, 28, 32, 36, 41, 49, 62, 63, 70, 77, 89, 114, 115, 117, 118, 142, 149, 158, 159, 182, 185, 186, 187, 192, 194, 196, 197, 198, 225, 292, 295, 298, 369, 371, 451, 459, 461, 508, 527, 540], [26, 28, 36, 37, 41, 49, 52, 62, 63, 65, 69, 77, 84, 109, 112, 114, 117, 118, 119, 142, 143, 145, 149, 158, 159, 182, 185, 186, 187, 192, 196, 197, 198, 199, 224, 298, 300, 307, 330, 369, 371, 379, 451, 459, 461, 508, 527, 537, 538, 540]], 'time-2': 147.17385935783386, 3: [[0.3953779935836792, 0.8903970122337341], [0.37462368607521057, 0.9063454270362854], [0.47964876890182495, 0.8965049386024475], [0.47953000664711, 0.9066847562789917]], 'masks-3': [[28, 41, 49, 62, 63, 77, 91, 117, 118, 149, 159, 185, 209, 302, 371, 386, 450, 451, 458, 537], [28, 32, 36, 39, 41, 49, 62, 63, 89, 117, 142, 149, 156, 159, 185, 186, 187, 192, 196, 287, 295, 298, 300, 302, 317, 378, 451, 459, 526, 537], [9, 28, 32, 36, 41, 49, 62, 63, 65, 69, 77, 114, 117, 118, 142, 147, 149, 158, 159, 185, 186, 187, 192, 196, 197, 198, 224, 282, 295, 298, 300, 302, 358, 369, 371, 450, 451, 458, 526, 540], [32, 36, 49, 50, 52, 53, 62, 63, 64, 65, 77, 89, 117, 118, 119, 142, 145, 149, 152, 158, 159, 182, 185, 186, 187, 189, 192, 196, 197, 198, 199, 212, 230, 289, 298, 301, 302, 317, 339, 369, 371, 378, 433, 450, 451, 459, 508, 527, 537, 540]], 'time-3': 150.6622612476349, 4: [[0.35613229870796204, 0.8920936584472656], [0.5296115279197693, 0.8887003660202026], [0.455687016248703, 0.9073634147644043], [0.46584126353263855, 0.910417377948761]], 'masks-4': [[28, 41, 63, 69, 77, 106, 117, 118, 149, 159, 185, 198, 268, 371, 450, 451, 457, 459, 537, 558], [28, 32, 41, 49, 62, 63, 70, 77, 106, 114, 117, 118, 142, 149, 159, 185, 186, 192, 198, 225, 298, 302, 342, 371, 451, 457, 459, 526, 537, 540], [28, 32, 36, 41, 52, 53, 62, 63, 70, 77, 89, 102, 114, 117, 118, 119, 142, 149, 159, 185, 186, 187, 192, 196, 197, 198, 212, 298, 302, 369, 371, 392, 450, 451, 457, 458, 461, 527, 537, 540], [28, 32, 36, 39, 40, 41, 62, 63, 69, 77, 108, 112, 114, 117, 118, 142, 149, 157, 158, 159, 185, 186, 187, 192, 196, 197, 198, 199, 224, 236, 274, 295, 298, 317, 323, 369, 371, 378, 380, 450, 451, 457, 458, 461, 489, 508, 526, 537, 540, 550]], 'time-4': 145.18028116226196}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6MNhfYqk5873"
      },
      "source": [
        "### Simple SA 60\n",
        "Sequential attention on 60% of the dataset (randomly sampled)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OcsfdNgq5873"
      },
      "outputs": [],
      "source": [
        "def get_model_simple_sa60(to_select):\n",
        "  return get_model_simple_sa(to_select)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "78XUs20N5873"
      },
      "outputs": [],
      "source": [
        "def get_sparse_model_evaluation_simple_sa60(mask):\n",
        "  return get_sparse_model_evaluation_simple_sa(mask)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5O0oJpwa5874"
      },
      "outputs": [],
      "source": [
        "def get_masks_fn_simple_sa60():\n",
        "  return get_masks_fn_simple_sa()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JNWkXwub5874"
      },
      "outputs": [],
      "source": [
        "# Reload global variables.\n",
        "(\n",
        "    ds_train,\n",
        "    x_train,\n",
        "    ds_val,\n",
        "    ds_test,\n",
        "    is_classification,\n",
        "    num_classes,\n",
        "    num_features,\n",
        "    num_train_steps_select,\n",
        "    loss_fn,\n",
        "    num_epochs_select,\n",
        ") = load_dataset(\n",
        "    num_epochs_select=624,\n",
        "    batch_size=256,\n",
        "    dataset='activity',\n",
        "    subsample_size=0.6,\n",
        ")\n",
        "simple_sa_60_experiment_data = run_experiment(\n",
        "    get_masks_fn_simple_sa60, get_sparse_model_evaluation_simple_sa60\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nfVCeecd5874"
      },
      "outputs": [],
      "source": [
        "print(simple_sa_60_experiment_data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "i5UZD0-R5874"
      },
      "outputs": [],
      "source": [
        "simple_sa_60_experiment_data = {0: [[0.4067540168762207, 0.885646402835846], [0.36136242747306824, 0.9151679873466492], [0.4017554223537445, 0.9110960364341736], [0.40648597478866577, 0.9250084757804871]], 'masks-0': [[28, 36, 50, 52, 62, 65, 91, 117, 142, 149, 159, 185, 192, 198, 281, 295, 450, 451, 459, 550], [28, 36, 39, 40, 50, 62, 63, 70, 112, 117, 142, 149, 159, 185, 186, 192, 197, 198, 229, 295, 298, 302, 371, 450, 451, 457, 459, 508, 527, 553], [28, 36, 38, 39, 41, 49, 62, 63, 65, 71, 77, 89, 114, 117, 118, 142, 145, 149, 156, 157, 158, 159, 182, 185, 186, 187, 192, 196, 197, 237, 278, 295, 302, 371, 374, 449, 450, 451, 459, 526], [28, 32, 36, 37, 38, 40, 53, 55, 62, 63, 65, 89, 106, 112, 117, 142, 145, 149, 156, 157, 159, 176, 182, 185, 186, 187, 189, 190, 192, 196, 197, 198, 199, 295, 297, 302, 307, 371, 374, 378, 447, 449, 450, 451, 457, 458, 508, 527, 537, 553]], 'time-0': 182.01329350471497, 1: [[0.4306696653366089, 0.8802171945571899], [0.3747563064098358, 0.9144893288612366], [0.4159783720970154, 0.9209365248680115], [0.39749303460121155, 0.9267051219940186]], 'masks-1': [[36, 40, 41, 60, 62, 91, 117, 142, 149, 159, 185, 192, 198, 302, 374, 450, 451, 457, 459, 553], [28, 36, 37, 39, 49, 53, 62, 63, 65, 69, 89, 117, 142, 149, 159, 182, 185, 186, 187, 192, 198, 302, 371, 374, 450, 451, 456, 459, 511, 527], [28, 36, 37, 38, 39, 41, 52, 55, 62, 63, 65, 70, 112, 114, 117, 118, 135, 142, 149, 182, 185, 186, 187, 192, 196, 197, 198, 212, 291, 295, 298, 369, 371, 374, 450, 451, 457, 458, 508, 553], [28, 32, 36, 37, 38, 39, 49, 50, 62, 63, 67, 69, 77, 89, 91, 103, 112, 114, 115, 117, 118, 142, 149, 156, 157, 158, 159, 182, 185, 186, 187, 192, 196, 197, 198, 199, 295, 299, 302, 371, 374, 449, 450, 451, 456, 459, 467, 514, 526, 553]], 'time-1': 176.2069034576416, 2: [[0.36632394790649414, 0.8903970122337341], [0.43216681480407715, 0.8975229263305664], [0.4646317660808563, 0.9168646335601807], [0.45503363013267517, 0.9100780487060547]], 'masks-2': [[28, 36, 41, 49, 62, 63, 65, 91, 117, 142, 149, 159, 185, 192, 281, 295, 378, 450, 451, 527], [16, 28, 36, 40, 53, 55, 63, 65, 69, 89, 112, 117, 142, 149, 159, 182, 185, 186, 187, 192, 197, 198, 298, 371, 374, 450, 451, 452, 459, 526], [28, 36, 50, 56, 62, 63, 71, 77, 112, 114, 115, 117, 118, 142, 149, 156, 158, 159, 182, 185, 186, 187, 192, 196, 197, 198, 282, 295, 298, 302, 371, 374, 378, 449, 450, 451, 456, 459, 508, 527], [28, 32, 36, 37, 38, 50, 52, 55, 62, 63, 66, 73, 76, 105, 112, 114, 115, 117, 118, 142, 145, 149, 152, 156, 157, 158, 159, 176, 182, 185, 186, 187, 192, 196, 197, 198, 209, 237, 298, 302, 374, 379, 450, 451, 452, 456, 459, 508, 518, 553]], 'time-2': 188.10492658615112, 3: [[0.3410423994064331, 0.903630793094635], [0.3846805989742279, 0.9043094515800476], [0.3857690989971161, 0.9148286581039429], [0.45191270112991333, 0.9090600609779358]], 'masks-3': [[36, 41, 52, 62, 63, 66, 91, 117, 142, 149, 159, 185, 198, 295, 302, 450, 451, 457, 458, 553], [22, 28, 36, 40, 53, 62, 63, 89, 112, 114, 117, 118, 142, 149, 157, 158, 159, 185, 186, 192, 196, 197, 302, 371, 374, 450, 451, 458, 527, 538], [28, 32, 36, 37, 38, 41, 49, 55, 62, 63, 65, 69, 112, 114, 117, 142, 149, 156, 157, 159, 182, 185, 186, 187, 192, 197, 198, 240, 264, 298, 302, 371, 374, 450, 451, 457, 458, 489, 537, 553], [28, 32, 36, 38, 39, 41, 52, 55, 62, 63, 74, 112, 114, 117, 118, 142, 145, 149, 152, 156, 157, 158, 159, 185, 186, 187, 192, 196, 197, 198, 236, 295, 297, 298, 302, 319, 371, 374, 378, 386, 423, 449, 450, 451, 457, 459, 508, 526, 538, 553]], 'time-3': 173.81576919555664, 4: [[0.4019664525985718, 0.884628415107727], [0.3444601893424988, 0.9175432920455933], [0.48131588101387024, 0.9032914638519287], [0.40348467230796814, 0.9144893288612366]], 'masks-4': [[28, 40, 41, 62, 74, 117, 142, 149, 159, 185, 186, 198, 302, 374, 450, 451, 457, 459, 529, 553], [28, 36, 52, 53, 62, 63, 65, 97, 112, 114, 117, 142, 149, 157, 159, 185, 186, 187, 192, 196, 197, 198, 295, 297, 374, 451, 457, 459, 508, 526], [25, 28, 32, 36, 39, 52, 53, 55, 62, 63, 70, 114, 117, 142, 149, 156, 158, 159, 182, 185, 186, 192, 196, 197, 198, 298, 302, 369, 371, 374, 386, 450, 451, 452, 456, 459, 461, 508, 527, 537], [15, 28, 32, 36, 39, 40, 50, 53, 62, 63, 65, 69, 89, 112, 114, 117, 142, 149, 157, 158, 159, 182, 185, 186, 187, 192, 196, 197, 198, 199, 295, 298, 302, 307, 369, 371, 374, 378, 386, 396, 450, 451, 458, 461, 476, 508, 527, 537, 552, 558]], 'time-4': 187.3709056377411}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mS2qSpNg5875"
      },
      "source": [
        "### Intermediate SA\n",
        "In comparison to Simple SA, Intermediate SA has a pool of K workers, that each compute the next best N features. At each step, the ensemble votes for a common mask, and all proceed with the same prior. In a setup where we need masks for [10, 20, 30] features, the ensemble first votes for 10 features, then with the common prior the ensemble votes again for the next 10 features, and so on until completion."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PBc9hgTz5875"
      },
      "outputs": [],
      "source": [
        "def get_model_intermediate_sa(\n",
        "    preselected_indices, to_select, num_train_steps_select\n",
        "):\n",
        "  model = SequentialAttentionMlp(\n",
        "      num_train_steps=num_train_steps_select,\n",
        "      num_candidates_to_select=to_select,\n",
        "      num_classes=num_classes,\n",
        "      layer_sequence=[67],\n",
        "      learning_rate=1e-5,\n",
        "      decay_steps=500,\n",
        "      decay_rate=1.0,\n",
        "      num_candidates=num_features,\n",
        "  )\n",
        "\n",
        "  model.compile(\n",
        "      loss=tf.keras.losses.CategoricalCrossentropy(), metrics=[\"accuracy\"]\n",
        "  )\n",
        "\n",
        "  return model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Bn6aihvt5875"
      },
      "outputs": [],
      "source": [
        "def get_sparse_model_evaluation_intermediate_sa(mask):\n",
        "  mlp_fit = SparseModel(\n",
        "      selected_features=tf.math.reduce_sum(\n",
        "          tf.one_hot(mask, num_features, dtype=tf.float32), 0\n",
        "      ),\n",
        "      num_classes=num_classes,\n",
        "      is_classification=True,\n",
        "      layer_sequence=[67],\n",
        "  )\n",
        "  mlp_fit.compile(loss=loss_fn, metrics=[\"accuracy\"])\n",
        "  mlp_fit.fit(\n",
        "      ds_train, validation_data=ds_val, epochs=num_epochs_fit, verbose=None\n",
        "  )\n",
        "  return mlp_fit.evaluate(ds_test)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NvZeU6m25875"
      },
      "outputs": [],
      "source": [
        "def get_masks_fn_intermediate_sa(num_epochs_select):\n",
        "  # parameters adapted for the constants defined in feature_counts\n",
        "  pool_size = 5\n",
        "  num_features_at_each_step = 20  # the gap for the features is at most 20; so at each step, we aim to fill preselected_features to the desired capacity\n",
        "  epochs_for_each = num_epochs_select\n",
        "\n",
        "  preselected_features = []\n",
        "\n",
        "  for feature_target in feature_counts:\n",
        "    print(f\"\\_ Getting mask for {feature_target} features.\")\n",
        "    model_pool = [\n",
        "        get_model_intermediate_sa(\n",
        "            preselected_features,\n",
        "            feature_target,\n",
        "            epochs_for_each * len(ds_train),\n",
        "        )\n",
        "        for i in range(pool_size)\n",
        "    ]\n",
        "\n",
        "    [\n",
        "        model.fit(\n",
        "            ds_train,\n",
        "            validation_data=ds_val,\n",
        "            epochs=epochs_for_each,\n",
        "            verbose=None,\n",
        "        )\n",
        "        for model in model_pool\n",
        "    ]\n",
        "\n",
        "    votes = {}\n",
        "\n",
        "    for model in model_pool:\n",
        "      selected_features = model.sa.selected_features\n",
        "      _, selected_indices = tf.math.top_k(selected_features, k=feature_target)\n",
        "      selected_indices = selected_indices.numpy()\n",
        "      for feature in list(selected_indices):\n",
        "        if feature not in preselected_features:\n",
        "          votes[feature] = votes.get(feature, 0) + 1\n",
        "\n",
        "    while len(preselected_features) \u003c feature_target:\n",
        "      max_feature = max(votes, key=votes.get)\n",
        "      votes[max_feature] = 0\n",
        "      preselected_features.append(max_feature)\n",
        "\n",
        "    print(f\"Features: {preselected_features}\")\n",
        "\n",
        "  return [\n",
        "      preselected_features[: feature_target + 1]\n",
        "      for feature_target in feature_counts\n",
        "  ]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "scxmOM6y5876"
      },
      "outputs": [],
      "source": [
        "# Reload global variables.\n",
        "(\n",
        "    ds_train,\n",
        "    x_train,\n",
        "    ds_val,\n",
        "    ds_test,\n",
        "    is_classification,\n",
        "    num_classes,\n",
        "    num_features,\n",
        "    num_train_steps_select,\n",
        "    loss_fn,\n",
        "    num_epochs_select,\n",
        ") = load_dataset(num_epochs_select=624, batch_size=256, dataset='activity')\n",
        "intermediate_sa_experiment_data = run_experiment(\n",
        "    lambda: get_masks_fn_intermediate_sa(num_epochs_select), get_sparse_model_evaluation_intermediate_sa\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vB4C0Ij05876"
      },
      "outputs": [],
      "source": [
        "print(intermediate_sa_experiment_data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jtmgOfHB5876"
      },
      "outputs": [],
      "source": [
        "intermediate_sa_experiment_data = {0: [[0.43645110726356506, 0.8832710981369019], [0.553971529006958, 0.8934509754180908], [0.5377596616744995, 0.9117746949195862], [0.46195361018180847, 0.926026463508606]], 'masks-0': [[37, 50, 89, 149, 159, 185, 198, 451, 28, 62, 69, 302, 458, 39, 40, 52, 182, 527, 49, 369, 186], [37, 50, 89, 149, 159, 185, 198, 451, 28, 62, 69, 302, 458, 39, 40, 52, 182, 527, 49, 369, 186, 36, 77, 142, 374, 41, 187, 192, 457, 55, 32], [37, 50, 89, 149, 159, 185, 198, 451, 28, 62, 69, 302, 458, 39, 40, 52, 182, 527, 49, 369, 186, 36, 77, 142, 374, 41, 187, 192, 457, 55, 32, 117, 450, 197, 298, 63, 371, 112, 295, 456, 461], [37, 50, 89, 149, 159, 185, 198, 451, 28, 62, 69, 302, 458, 39, 40, 52, 182, 527, 49, 369, 186, 36, 77, 142, 374, 41, 187, 192, 457, 55, 32, 117, 450, 197, 298, 63, 371, 112, 295, 456, 461, 118, 459, 508, 526, 537, 145, 157, 196, 553]], 'time-0': 1264.768618106842, 1: [[0.3338741362094879, 0.9032914638519287], [0.5057540535926819, 0.8958262801170349], [0.40333858132362366, 0.92840176820755], [0.37046849727630615, 0.9443501830101013]], 'masks-1': [[37, 39, 149, 159, 198, 451, 89, 185, 450, 40, 50, 52, 69, 459, 62, 302, 65, 374, 70, 371, 142], [37, 39, 149, 159, 198, 451, 89, 185, 450, 40, 50, 52, 69, 459, 62, 302, 65, 374, 70, 371, 142, 186, 36, 117, 197, 369, 28, 41, 187, 192, 32], [37, 39, 149, 159, 198, 451, 89, 185, 450, 40, 50, 52, 69, 459, 62, 302, 65, 374, 70, 371, 142, 186, 36, 117, 197, 369, 28, 41, 187, 192, 32, 63, 457, 77, 112, 508, 295, 553, 145, 158, 182], [37, 39, 149, 159, 198, 451, 89, 185, 450, 40, 50, 52, 69, 459, 62, 302, 65, 374, 70, 371, 142, 186, 36, 117, 197, 369, 28, 41, 187, 192, 32, 63, 457, 77, 112, 508, 295, 553, 145, 158, 182, 526, 114, 298, 456, 461, 57, 199, 452, 157]], 'time-1': 1288.7287442684174, 2: [[0.29957863688468933, 0.9114353656768799], [0.4075995683670044, 0.913810670375824], [0.4020655155181885, 0.9263657927513123], [0.4420371353626251, 0.9175432920455933]], 'masks-2': [[37, 149, 159, 198, 450, 451, 39, 50, 69, 185, 369, 62, 456, 28, 459, 52, 142, 374, 458, 508, 89], [37, 149, 159, 198, 450, 451, 39, 50, 69, 185, 369, 62, 456, 28, 459, 52, 142, 374, 458, 508, 89, 186, 192, 298, 302, 371, 49, 457, 77, 182, 187], [37, 149, 159, 198, 450, 451, 39, 50, 69, 185, 369, 62, 456, 28, 459, 52, 142, 374, 458, 508, 89, 186, 192, 298, 302, 371, 49, 457, 77, 182, 187, 32, 36, 117, 295, 526, 158, 197, 300, 40, 112], [37, 149, 159, 198, 450, 451, 39, 50, 69, 185, 369, 62, 456, 28, 459, 52, 142, 374, 458, 508, 89, 186, 192, 298, 302, 371, 49, 457, 77, 182, 187, 32, 36, 117, 295, 526, 158, 197, 300, 40, 112, 114, 553, 63, 41, 118, 157, 296, 461, 527]], 'time-2': 1285.4171187877655, 3: [[0.4555092453956604, 0.879199206829071], [0.335846483707428, 0.913810670375824], [0.43378233909606934, 0.9239904880523682], [0.3180316090583801, 0.9426535367965698]], 'masks-3': [[149, 159, 185, 198, 451, 37, 62, 89, 369, 450, 39, 50, 374, 69, 302, 459, 295, 142, 186, 70, 41], [149, 159, 185, 198, 451, 37, 62, 89, 369, 450, 39, 50, 374, 69, 302, 459, 295, 142, 186, 70, 41, 117, 182, 49, 32, 63, 91, 187, 36, 298, 28], [149, 159, 185, 198, 451, 37, 62, 89, 369, 450, 39, 50, 374, 69, 302, 459, 295, 142, 186, 70, 41, 117, 182, 49, 32, 63, 91, 187, 36, 298, 28, 192, 371, 457, 65, 77, 112, 300, 526, 158, 197], [149, 159, 185, 198, 451, 37, 62, 89, 369, 450, 39, 50, 374, 69, 302, 459, 295, 142, 186, 70, 41, 117, 182, 49, 32, 63, 91, 187, 36, 298, 28, 192, 371, 457, 65, 77, 112, 300, 526, 158, 197, 458, 145, 196, 53, 57, 456, 461, 508, 537]], 'time-3': 1288.1895353794098, 4: [[0.35450759530067444, 0.9022734761238098], [0.5208109617233276, 0.894808292388916], [0.3706172704696655, 0.9253478050231934], [0.40247786045074463, 0.9304377436637878]], 'masks-4': [[37, 149, 185, 198, 374, 451, 50, 69, 159, 369, 459, 39, 40, 62, 302, 456, 28, 89, 371, 450, 142], [37, 149, 185, 198, 374, 451, 50, 69, 159, 369, 459, 39, 40, 62, 302, 456, 28, 89, 371, 450, 142, 186, 36, 117, 187, 295, 49, 41, 192, 508, 32], [37, 149, 185, 198, 374, 451, 50, 69, 159, 369, 459, 39, 40, 62, 302, 456, 28, 89, 371, 450, 142, 186, 36, 117, 187, 295, 49, 41, 192, 508, 32, 77, 182, 197, 63, 65, 112, 145, 298, 452, 458], [37, 149, 185, 198, 374, 451, 50, 69, 159, 369, 459, 39, 40, 62, 302, 456, 28, 89, 371, 450, 142, 186, 36, 117, 187, 295, 49, 41, 192, 508, 32, 77, 182, 197, 63, 65, 112, 145, 298, 452, 458, 114, 461, 537, 156, 157, 158, 527, 52, 199]], 'time-4': 1246.1731293201447}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7ZJRM-_o5876"
      },
      "source": [
        "### LLYMask"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Fw4GsNBH5876"
      },
      "outputs": [],
      "source": [
        "def get_model_lly(to_select):\n",
        "  model = LiaoLattyYangModel(\n",
        "      num_classes=num_classes,\n",
        "      layer_sequence=[67],\n",
        "      num_inputs=num_features,\n",
        "  )\n",
        "\n",
        "  model.compile(\n",
        "      loss=tf.keras.losses.CategoricalCrossentropy(), metrics=[\"accuracy\"]\n",
        "  )\n",
        "\n",
        "  return model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N4633pd15876"
      },
      "outputs": [],
      "source": [
        "def get_sparse_model_evaluation_lly(mask):\n",
        "  mlp_fit = SparseModel(\n",
        "      selected_features=tf.math.reduce_sum(\n",
        "          tf.one_hot(mask, num_features, dtype=tf.float32), 0\n",
        "      ),\n",
        "      num_classes=num_classes,\n",
        "      is_classification=True,\n",
        "      layer_sequence=[67],\n",
        "  )\n",
        "  mlp_fit.compile(loss=loss_fn, metrics=[\"accuracy\"])\n",
        "  mlp_fit.fit(\n",
        "      ds_train, validation_data=ds_val, epochs=num_epochs_fit, verbose=None\n",
        "  )\n",
        "  return mlp_fit.evaluate(ds_test)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dHABRlvT5877"
      },
      "outputs": [],
      "source": [
        "def get_masks_fn_lly():\n",
        "  masks = []\n",
        "  for feature_count in feature_counts:\n",
        "    print(f\"\\_ Getting mask for {feature_count} features.\")\n",
        "    mask = []\n",
        "\n",
        "    model = get_model_lly(feature_count)\n",
        "    model.fit(\n",
        "        ds_train, validation_data=ds_val, epochs=num_epochs_select, verbose=None\n",
        "    )\n",
        "\n",
        "    attention_logits = model.lly(tf.convert_to_tensor(x_train))\n",
        "    _, selected_indices = tf.math.top_k(attention_logits, k=feature_count)\n",
        "    selected_indices = selected_indices.numpy()\n",
        "\n",
        "    masks.append(list(selected_indices))\n",
        "\n",
        "  return masks"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HNLL01g-5877"
      },
      "outputs": [],
      "source": [
        "# Reload global variables.\n",
        "(\n",
        "    ds_train,\n",
        "    x_train,\n",
        "    ds_val,\n",
        "    ds_test,\n",
        "    is_classification,\n",
        "    num_classes,\n",
        "    num_features,\n",
        "    num_train_steps_select,\n",
        "    loss_fn,\n",
        "    num_epochs_select,\n",
        ") = load_dataset(num_epochs_select=624, batch_size=256, dataset='activity')\n",
        "lly_experiment_data = run_experiment(\n",
        "    get_masks_fn_lly, get_sparse_model_evaluation_lly\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m5EYeenJ5877"
      },
      "outputs": [],
      "source": [
        "print(lly_experiment_data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nJkatacu5877"
      },
      "outputs": [],
      "source": [
        "lly_experiment_data = {0: [[0.7163563370704651, 0.678656280040741], [0.5225291848182678, 0.8177807927131653], [0.24569784104824066, 0.9209365248680115], [0.40209734439849854, 0.8435697555541992]], 'masks-0': [[368, 87, 503, 203, 352, 214, 281, 234, 272, 280, 180, 367, 181, 165, 18, 275, 245, 510, 134, 140], [268, 265, 85, 4, 520, 281, 9, 275, 25, 306, 7, 508, 126, 125, 443, 288, 57, 507, 90, 33, 528, 84, 547, 366, 522, 442, 163, 504, 272, 466], [522, 129, 226, 502, 272, 201, 19, 170, 360, 320, 252, 430, 100, 284, 102, 518, 103, 438, 260, 165, 50, 283, 125, 380, 27, 239, 247, 145, 541, 424, 254, 285, 375, 49, 92, 367, 23, 86, 193, 231], [201, 184, 347, 227, 181, 324, 123, 246, 504, 83, 543, 284, 390, 218, 288, 104, 259, 359, 228, 58, 105, 274, 100, 507, 125, 393, 518, 23, 255, 378, 265, 472, 140, 90, 18, 220, 89, 542, 474, 128, 528, 69, 182, 204, 85, 296, 171, 452, 13, 169]], 'time-0': 8.342309474945068, 1: [[0.7117600440979004, 0.7105531096458435], [0.5455102920532227, 0.7845266461372375], [0.423885315656662, 0.8805565237998962], [0.4768899381160736, 0.8751272559165955]], 'masks-1': [[445, 276, 408, 135, 346, 139, 439, 328, 442, 423, 247, 400, 37, 533, 138, 404, 413, 193, 98, 186], [93, 124, 523, 276, 410, 217, 517, 260, 518, 282, 25, 544, 171, 102, 129, 482, 428, 471, 167, 109, 372, 240, 96, 230, 507, 100, 30, 503, 287, 341], [200, 126, 193, 83, 533, 12, 165, 265, 285, 93, 347, 53, 432, 351, 3, 312, 109, 219, 273, 288, 242, 317, 339, 502, 545, 346, 503, 326, 85, 22, 88, 181, 336, 153, 372, 479, 379, 183, 103, 218], [200, 19, 40, 218, 256, 522, 228, 83, 204, 425, 267, 382, 155, 240, 57, 509, 545, 515, 342, 94, 213, 235, 216, 15, 501, 397, 89, 241, 140, 372, 50, 504, 543, 497, 349, 516, 470, 258, 391, 433, 10, 394, 217, 229, 365, 523, 62, 491, 227, 409]], 'time-1': 8.528566122055054, 2: [[0.4011997878551483, 0.8608754873275757], [0.33971574902534485, 0.8598574995994568], [0.4684300720691681, 0.8042076826095581], [0.4397214949131012, 0.891414999961853]], 'masks-2': [[271, 20, 425, 515, 345, 272, 40, 234, 5, 213, 53, 153, 325, 128, 243, 154, 403, 483, 361, 286], [179, 544, 21, 523, 546, 507, 58, 424, 509, 272, 345, 269, 84, 169, 207, 529, 233, 271, 438, 275, 23, 389, 85, 26, 40, 227, 444, 498, 541, 368], [84, 83, 407, 11, 366, 389, 256, 102, 356, 125, 282, 90, 432, 135, 202, 464, 372, 393, 179, 224, 360, 14, 421, 134, 427, 213, 503, 377, 364, 203, 20, 273, 269, 220, 353, 449, 216, 257, 444, 337], [182, 86, 446, 125, 205, 20, 93, 94, 520, 366, 217, 246, 99, 367, 40, 9, 181, 531, 485, 252, 549, 58, 42, 270, 332, 342, 165, 482, 496, 3, 428, 8, 397, 193, 191, 396, 476, 231, 63, 213, 257, 92, 183, 185, 501, 76, 308, 518, 41, 145]], 'time-2': 8.443706274032593, 3: [[0.6048743724822998, 0.8483203053474426], [0.5474429726600647, 0.7980997562408447], [0.26758962869644165, 0.8836104273796082], [0.5063556432723999, 0.8076009750366211]], 'masks-3': [[140, 429, 347, 184, 92, 533, 350, 143, 24, 133, 274, 531, 252, 471, 6, 559, 16, 272, 57, 123], [366, 271, 259, 7, 428, 504, 267, 474, 280, 431, 394, 13, 324, 284, 410, 265, 272, 200, 179, 125, 433, 260, 276, 103, 399, 4, 360, 472, 541, 204], [426, 267, 88, 367, 523, 244, 140, 260, 231, 549, 95, 393, 135, 442, 389, 286, 226, 266, 352, 134, 528, 522, 52, 40, 473, 239, 271, 504, 405, 358, 23, 288, 425, 403, 275, 351, 530, 518, 289, 346], [218, 442, 232, 284, 312, 256, 93, 510, 172, 429, 16, 347, 15, 229, 523, 23, 425, 90, 346, 5, 131, 165, 145, 433, 174, 4, 84, 33, 141, 153, 431, 500, 7, 287, 186, 76, 371, 535, 314, 154, 149, 14, 514, 262, 366, 362, 548, 498, 326, 116]], 'time-3': 8.61705470085144, 4: [[0.33473795652389526, 0.8652867078781128], [0.48786628246307373, 0.8649473786354065], [0.42214009165763855, 0.9015948176383972], [0.4944632649421692, 0.810654878616333]], 'masks-4': [[56, 260, 100, 346, 286, 359, 288, 40, 364, 230, 79, 444, 327, 179, 430, 16, 244, 426, 517, 522], [93, 102, 101, 220, 129, 503, 446, 130, 265, 174, 421, 267, 254, 490, 316, 544, 30, 344, 15, 430, 269, 542, 282, 494, 202, 41, 242, 322, 543, 534], [56, 208, 504, 281, 83, 113, 103, 177, 549, 104, 163, 443, 434, 269, 426, 515, 273, 50, 387, 324, 518, 25, 225, 55, 488, 14, 477, 379, 511, 446, 144, 439, 124, 529, 345, 406, 299, 338, 318, 238], [520, 324, 530, 254, 201, 181, 445, 504, 63, 95, 266, 304, 154, 104, 245, 361, 443, 127, 14, 146, 278, 103, 346, 382, 393, 217, 246, 183, 6, 193, 167, 308, 432, 509, 521, 510, 275, 239, 541, 334, 442, 407, 78, 318, 135, 386, 439, 280, 220, 100]], 'time-4': 8.696548700332642}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jjD1ZxD-5877"
      },
      "source": [
        "### Group Lasso"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3N4EzZSH5878"
      },
      "outputs": [],
      "source": [
        "def get_model_group_lasso(to_select):\n",
        "  model = SequentialLASSOModel(\n",
        "      num_features,\n",
        "      to_select,\n",
        "      num_train_steps=num_train_steps_select,\n",
        "      num_classes=num_classes,\n",
        "      layer_sequence=[67],\n",
        "      group_lasso_scale=0.01,\n",
        "  )\n",
        "\n",
        "  model.compile(\n",
        "      loss=tf.keras.losses.CategoricalCrossentropy(), metrics=[\"accuracy\"]\n",
        "  )\n",
        "\n",
        "  return model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "L5GKtrO55878"
      },
      "outputs": [],
      "source": [
        "def get_sparse_model_evaluation_group_lasso(mask):\n",
        "  mlp_fit = SparseModel(\n",
        "      selected_features=tf.math.reduce_sum(\n",
        "          tf.one_hot(mask, num_features, dtype=tf.float32), 0\n",
        "      ),\n",
        "      num_classes=num_classes,\n",
        "      is_classification=True,\n",
        "      layer_sequence=[67],\n",
        "  )\n",
        "  mlp_fit.compile(loss=loss_fn, metrics=[\"accuracy\"])\n",
        "  mlp_fit.fit(\n",
        "      ds_train, validation_data=ds_val, epochs=num_epochs_fit, verbose=None\n",
        "  )\n",
        "  return mlp_fit.evaluate(ds_test)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gCSeHJrN5878"
      },
      "outputs": [],
      "source": [
        "def get_masks_fn_group_lasso():\n",
        "  masks = []\n",
        "  for feature_count in feature_counts:\n",
        "    print(f\"\\_ Getting mask for {feature_count} features.\")\n",
        "    mask = []\n",
        "\n",
        "    model = get_model_group_lasso(feature_count)\n",
        "    model.fit(\n",
        "        ds_train, validation_data=ds_val, epochs=num_epochs_select, verbose=None\n",
        "    )\n",
        "\n",
        "    selected_indices = model.seql.selected_features_history.numpy().tolist()\n",
        "\n",
        "    masks.append(list(selected_indices))\n",
        "\n",
        "  return masks"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Xai67aSb5878"
      },
      "outputs": [],
      "source": [
        "# Reload global variables.\n",
        "(\n",
        "    ds_train,\n",
        "    x_train,\n",
        "    ds_val,\n",
        "    ds_test,\n",
        "    is_classification,\n",
        "    num_classes,\n",
        "    num_features,\n",
        "    num_train_steps_select,\n",
        "    loss_fn,\n",
        "    num_epochs_select,\n",
        ") = load_dataset(num_epochs_select=624, batch_size=256, dataset='activity')\n",
        "group_lasso_experiment_data = run_experiment(\n",
        "    get_masks_fn_group_lasso, get_sparse_model_evaluation_group_lasso\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TDK_Mt125878"
      },
      "outputs": [],
      "source": [
        "print(group_lasso_experiment_data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7aSpmJIp5879"
      },
      "outputs": [],
      "source": [
        "group_lasso_experiment_data = {0: [[0.4226674735546112, 0.91923987865448], [0.634773850440979, 0.9032914638519287], [0.5911436676979065, 0.9168646335601807], [0.6113049983978271, 0.9155073165893555]], 'masks-0': [[50, 182, 52, 104, 508, 37, 186, 185, 149, 69, 433, 179, 459, 198, 197, 374, 526, 36, 300, 187], [50, 182, 508, 52, 104, 37, 186, 185, 149, 433, 69, 459, 461, 197, 198, 526, 374, 187, 36, 114, 156, 194, 118, 65, 297, 82, 112, 158, 190, 1], [50, 182, 508, 52, 104, 37, 186, 149, 185, 433, 69, 461, 459, 198, 374, 197, 187, 36, 526, 114, 158, 119, 194, 65, 156, 112, 196, 82, 152, 300, 118, 190, 38, 557, 63, 142, 298, 32, 370, 290], [50, 182, 508, 52, 37, 104, 186, 149, 185, 433, 69, 459, 461, 198, 197, 374, 526, 187, 36, 114, 194, 156, 112, 119, 118, 158, 196, 65, 190, 300, 152, 82, 557, 552, 298, 297, 32, 142, 78, 295, 451, 55, 263, 551, 238, 38, 448, 80, 192, 370]], 'time-0': 226.19938397407532, 1: [[0.3658556640148163, 0.9209365248680115], [0.5628268122673035, 0.9195792078971863], [0.5405339002609253, 0.922633171081543], [0.5470306277275085, 0.9175432920455933]], 'masks-1': [[508, 182, 52, 50, 104, 37, 186, 149, 185, 433, 69, 461, 459, 187, 526, 411, 197, 198, 62, 36], [50, 182, 508, 52, 104, 37, 186, 185, 149, 433, 69, 459, 179, 198, 197, 374, 36, 526, 187, 65, 112, 114, 82, 300, 158, 118, 513, 156, 194, 119], [50, 182, 508, 52, 104, 37, 186, 185, 149, 433, 69, 374, 179, 198, 459, 197, 187, 36, 526, 114, 63, 369, 300, 119, 112, 158, 194, 82, 196, 156, 190, 152, 118, 65, 557, 552, 298, 32, 38, 55], [50, 182, 508, 52, 37, 104, 186, 149, 185, 433, 102, 69, 459, 198, 65, 197, 187, 374, 36, 190, 114, 112, 526, 300, 118, 156, 158, 194, 196, 32, 179, 557, 119, 298, 152, 370, 142, 552, 238, 80, 236, 38, 451, 63, 263, 448, 291, 551, 295, 82]], 'time-1': 221.96721386909485, 2: [[0.38084718585014343, 0.9256871342658997], [0.5515010952949524, 0.9161859750747681], [0.5416906476020813, 0.922633171081543], [0.5043313503265381, 0.9239904880523682]], 'masks-2': [[508, 182, 52, 50, 104, 102, 37, 186, 149, 185, 69, 433, 198, 459, 65, 374, 36, 187, 526, 513], [50, 182, 508, 52, 104, 37, 186, 185, 149, 69, 433, 459, 461, 198, 197, 526, 187, 112, 411, 62, 119, 374, 236, 82, 36, 158, 369, 196, 55, 118], [50, 182, 508, 52, 37, 104, 186, 102, 185, 149, 433, 69, 459, 198, 197, 374, 187, 114, 36, 526, 65, 112, 513, 158, 156, 194, 190, 196, 300, 119, 32, 370, 179, 557, 142, 552, 298, 118, 38, 238], [50, 182, 508, 52, 37, 104, 186, 149, 185, 433, 69, 102, 459, 198, 461, 197, 526, 247, 36, 187, 374, 411, 65, 112, 118, 114, 158, 142, 119, 156, 194, 300, 557, 152, 196, 32, 38, 190, 78, 236, 238, 55, 448, 370, 451, 450, 121, 298, 263, 550]], 'time-2': 220.07309317588806, 3: [[0.38559845089912415, 0.9202578663825989], [0.5503714084625244, 0.9165253043174744], [0.601382851600647, 0.9175432920455933], [0.5935981273651123, 0.913810670375824]], 'masks-3': [[50, 182, 52, 104, 508, 37, 186, 149, 185, 102, 69, 459, 65, 198, 187, 197, 179, 526, 36, 300], [50, 182, 508, 52, 104, 37, 186, 102, 149, 185, 69, 459, 433, 198, 65, 187, 36, 526, 374, 114, 197, 158, 297, 1, 118, 194, 156, 112, 82, 32], [508, 182, 50, 52, 37, 104, 186, 185, 433, 247, 149, 69, 461, 198, 197, 459, 526, 187, 374, 36, 114, 156, 118, 65, 194, 142, 121, 300, 158, 1, 55, 38, 557, 119, 82, 448, 236, 112, 152, 32], [50, 182, 508, 52, 37, 104, 186, 433, 185, 149, 69, 459, 461, 198, 197, 374, 187, 526, 36, 114, 62, 112, 28, 190, 236, 196, 194, 156, 300, 118, 158, 152, 557, 142, 65, 298, 82, 119, 238, 552, 38, 32, 295, 451, 55, 370, 448, 2, 263, 290]], 'time-3': 216.88831615447998, 4: [[0.3887333571910858, 0.9250084757804871], [0.5634514689445496, 0.9158466458320618], [0.6331954598426819, 0.9155073165893555], [0.55201655626297, 0.9178826212882996]], 'masks-4': [[182, 508, 52, 50, 104, 37, 186, 102, 149, 247, 69, 145, 187, 461, 198, 526, 459, 197, 36, 374], [50, 182, 508, 52, 104, 37, 186, 185, 149, 433, 69, 459, 179, 197, 198, 526, 36, 187, 374, 65, 114, 300, 112, 119, 196, 156, 194, 82, 152, 118], [50, 182, 508, 52, 104, 37, 186, 185, 149, 433, 102, 69, 461, 198, 459, 187, 197, 526, 36, 374, 119, 114, 194, 156, 112, 158, 118, 65, 196, 300, 152, 82, 190, 80, 557, 32, 552, 142, 298, 38], [182, 508, 50, 52, 37, 104, 186, 433, 185, 149, 69, 459, 461, 197, 198, 374, 187, 36, 526, 119, 114, 158, 194, 156, 196, 112, 65, 190, 300, 82, 152, 557, 298, 118, 142, 78, 32, 55, 551, 63, 370, 451, 263, 295, 38, 552, 238, 236, 369, 290]], 'time-4': 213.59599232673645}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7HpiTRwLZqK-"
      },
      "source": [
        "## Plots (Activity)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nd549yINZsRV"
      },
      "outputs": [],
      "source": [
        "experiments_data = [\n",
        "    get_data_from_experiment_output(simple_sa_experiment_data),\n",
        "    get_data_from_experiment_output(simple_sa_20_experiment_data),\n",
        "    get_data_from_experiment_output(simple_sa_40_experiment_data),\n",
        "    get_data_from_experiment_output(simple_sa_60_experiment_data),\n",
        "    get_data_from_experiment_output(\n",
        "        intermediate_sa_experiment_data, divide_time_by=5\n",
        "    ),  # ensemble of 5 models that can run in parallel\n",
        "    get_data_from_experiment_output(lly_experiment_data),\n",
        "    get_data_from_experiment_output(group_lasso_experiment_data),\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1Xtjk6zgZsRX"
      },
      "outputs": [],
      "source": [
        "fig = plt.gcf()\n",
        "fig.set_size_inches(18.5, 12.5)\n",
        "\n",
        "colors_for_feature_count = ['b', 'm', 'y', 'g']\n",
        "\n",
        "method_x = []\n",
        "method_y_mean = []\n",
        "method_y_std = []\n",
        "colors = []\n",
        "\n",
        "for feature_count_i in range(len(feature_counts)):\n",
        "  feature_count = feature_counts[feature_count_i]\n",
        "\n",
        "  for i in range(len(experiments_data)):\n",
        "    method = methods[i]\n",
        "    experiment_data = experiments_data[i]\n",
        "\n",
        "    method_x.append(f'{method} ({feature_count}f)')\n",
        "    method_y_mean.append(experiment_data['mean_acc_for_mask'][feature_count])\n",
        "    method_y_std.append(experiment_data['std_acc_for_mask'][feature_count])\n",
        "    colors.append(colors_for_feature_count[feature_count_i])\n",
        "\n",
        "method_x_pos = np.arange(len(method_x))\n",
        "\n",
        "fig, ax = plt.subplots()\n",
        "barlist = ax.bar(\n",
        "    method_x,\n",
        "    method_y_mean,\n",
        "    yerr=method_y_std,\n",
        "    align='center',\n",
        "    alpha=0.5,\n",
        "    ecolor='black',\n",
        "    capsize=10,\n",
        ")\n",
        "\n",
        "for i in range(len(barlist)):\n",
        "  barlist[i].set_color(colors[i])\n",
        "\n",
        "\n",
        "ax.set_xticks(method_x_pos)\n",
        "ax.set_xticklabels(method_x)\n",
        "ax.yaxis.grid(True)\n",
        "\n",
        "plt.ylabel('Accuracy')\n",
        "plt.xlabel('Method and feature set size')\n",
        "plt.xticks(rotation=30)\n",
        "ax.set_ylim([0.7, 0.98])\n",
        "plt.rcParams['figure.figsize'] = (22, 9)\n",
        "plt.title(\n",
        "    'Accuracy (mean and std) for different feature selection methods, with'\n",
        "    ' different feature set size targets, groupped by feature set size'\n",
        "    ' (Activity).'\n",
        ")\n",
        "\n",
        "red_patch = mpatches.Patch(color='red', label='The red data')\n",
        "blue_patch = mpatches.Patch(color='blue', label='The blue data')\n",
        "\n",
        "plt.legend(\n",
        "    handles=[\n",
        "        *[\n",
        "            mpatches.Patch(\n",
        "                color=color,\n",
        "                label=(\n",
        "                    f'{feature_count} features'\n",
        "                    f' ({round(feature_count / num_features *100, 2)}%)'\n",
        "                ),\n",
        "            )\n",
        "            for (color, feature_count) in zip(\n",
        "                colors_for_feature_count, feature_counts\n",
        "            )\n",
        "        ],\n",
        "        mpatches.Patch(\n",
        "            color='black',\n",
        "            label=(\n",
        "                'SimplaSA\u003cX\u003e: Sequential Attention on a uniform sample of X% of'\n",
        "                ' the datapoints.'\n",
        "            ),\n",
        "        ),\n",
        "    ]\n",
        ")\n",
        "\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W_94zjzRYLGT"
      },
      "outputs": [],
      "source": [
        "fig = plt.gcf()\n",
        "fig.set_size_inches(18.5, 12.5)\n",
        "\n",
        "colors_for_feature_count = ['b', 'm', 'y', 'g']\n",
        "\n",
        "method_x = []\n",
        "method_y_mean = []\n",
        "method_y_std = []\n",
        "colors = []\n",
        "\n",
        "for feature_count_i in range(len(feature_counts)):\n",
        "  feature_count = feature_counts[feature_count_i]\n",
        "\n",
        "  for i in range(len(experiments_data)):\n",
        "    method = methods[i]\n",
        "    experiment_data = experiments_data[i]\n",
        "\n",
        "    method_x.append(f'{method} ({feature_count}f)')\n",
        "    method_y_mean.append(experiment_data['mean_acc_for_mask'][feature_count])\n",
        "    method_y_std.append(experiment_data['std_acc_for_mask'][feature_count])\n",
        "    colors.append(colors_for_feature_count[feature_count_i])\n",
        "\n",
        "method_x_pos = np.arange(len(method_x))\n",
        "\n",
        "fig, ax = plt.subplots()\n",
        "barlist = ax.bar(\n",
        "    method_x,\n",
        "    method_y_mean,\n",
        "    yerr=method_y_std,\n",
        "    align='center',\n",
        "    alpha=0.5,\n",
        "    ecolor='black',\n",
        "    capsize=10,\n",
        ")\n",
        "\n",
        "for i in range(len(barlist)):\n",
        "  barlist[i].set_color(colors[i])\n",
        "\n",
        "\n",
        "ax.set_xticks(method_x_pos)\n",
        "ax.set_xticklabels(method_x)\n",
        "ax.yaxis.grid(True)\n",
        "\n",
        "plt.ylabel('Accuracy')\n",
        "plt.xlabel('Method and feature set size')\n",
        "plt.xticks(rotation=30)\n",
        "ax.set_ylim([0.7, 0.98])\n",
        "plt.rcParams['figure.figsize'] = (22, 9)\n",
        "plt.title(\n",
        "    'Accuracy (mean and std) for different feature selection methods, with'\n",
        "    ' different feature set size targets, groupped by feature set size'\n",
        "    ' (Activity).'\n",
        ")\n",
        "\n",
        "red_patch = mpatches.Patch(color='red', label='The red data')\n",
        "blue_patch = mpatches.Patch(color='blue', label='The blue data')\n",
        "\n",
        "plt.legend(\n",
        "    handles=[\n",
        "        *[\n",
        "            mpatches.Patch(\n",
        "                color=color,\n",
        "                label=(\n",
        "                    f'{feature_count} features'\n",
        "                    f' ({round(feature_count / num_features *100, 2)}%)'\n",
        "                ),\n",
        "            )\n",
        "            for (color, feature_count) in zip(\n",
        "                colors_for_feature_count, feature_counts\n",
        "            )\n",
        "        ],\n",
        "        mpatches.Patch(\n",
        "            color='black',\n",
        "            label=(\n",
        "                'SimplaSA\u003cX\u003e: Sequential Attention on a uniform sample of X% of'\n",
        "                ' the datapoints.'\n",
        "            ),\n",
        "        ),\n",
        "    ]\n",
        ")\n",
        "\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Iv5QIrRycvQg"
      },
      "outputs": [],
      "source": [
        "fig = plt.gcf()\n",
        "fig.set_size_inches(18.5, 12.5)\n",
        "\n",
        "method_x = []\n",
        "method_y_mean = []\n",
        "method_y_std = []\n",
        "\n",
        "for i in range(len(experiments_data)):\n",
        "  method = methods[i]\n",
        "  experiment_data = experiments_data[i]\n",
        "\n",
        "  method_x.append(method)\n",
        "  method_y_mean.append(experiment_data['mean_time'])\n",
        "  method_y_std.append(experiment_data['std_time'])\n",
        "\n",
        "method_x_pos = np.arange(len(method_x))\n",
        "\n",
        "fig, ax = plt.subplots()\n",
        "barlist = ax.bar(\n",
        "    method_x,\n",
        "    method_y_mean,\n",
        "    yerr=method_y_std,\n",
        "    align='center',\n",
        "    alpha=0.5,\n",
        "    ecolor='black',\n",
        "    capsize=10,\n",
        ")\n",
        "\n",
        "ax.set_xticks(method_x_pos)\n",
        "ax.set_xticklabels(method_x)\n",
        "ax.yaxis.grid(True)\n",
        "\n",
        "plt.ylabel('Seconds')\n",
        "plt.xlabel('Method')\n",
        "plt.rcParams['figure.figsize'] = (22, 9)\n",
        "plt.title(\n",
        "    'End-to-end time to compute masks for all different sizes (Activity).'\n",
        ")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "macFRd2_7Bz0"
      },
      "outputs": [],
      "source": [
        "divide_by = experiments_data[0]['mean_time']\n",
        "\n",
        "fig = plt.gcf()\n",
        "fig.set_size_inches(18.5, 12.5)\n",
        "\n",
        "method_x = []\n",
        "method_y_mean = []\n",
        "method_y_std = []\n",
        "\n",
        "for i in range(len(experiments_data)):\n",
        "  method = methods[i]\n",
        "  experiment_data = experiments_data[i]\n",
        "\n",
        "  method_x.append(method)\n",
        "  method_y_mean.append(experiment_data['mean_time'] / divide_by)\n",
        "  method_y_std.append(experiment_data['std_time'] / divide_by)\n",
        "\n",
        "method_x_pos = np.arange(len(method_x))\n",
        "\n",
        "fig, ax = plt.subplots()\n",
        "barlist = ax.bar(\n",
        "    method_x,\n",
        "    method_y_mean,\n",
        "    yerr=method_y_std,\n",
        "    align='center',\n",
        "    alpha=0.5,\n",
        "    ecolor='black',\n",
        "    capsize=10,\n",
        ")\n",
        "\n",
        "ax.set_xticks(method_x_pos)\n",
        "ax.set_xticklabels(method_x)\n",
        "ax.yaxis.grid(True)\n",
        "\n",
        "plt.ylabel('Time relative to simple Sequential Attention')\n",
        "plt.xlabel('Method')\n",
        "plt.rcParams['figure.figsize'] = (22, 9)\n",
        "plt.title('Time to compute masks for all different sizes (Activity).')\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ovzFxceLLMmB"
      },
      "outputs": [],
      "source": [
        "divide_by = (\n",
        "    experiment_data['mean_acc_for_mask'][feature_counts[-1]]\n",
        "    / experiment_data['mean_time']\n",
        ")\n",
        "fig = plt.gcf()\n",
        "fig.set_size_inches(18.5, 12.5)\n",
        "\n",
        "method_x = []\n",
        "method_y_mean = []\n",
        "method_y_std = []\n",
        "\n",
        "for i in range(len(experiments_data)):\n",
        "  method = methods[i]\n",
        "  experiment_data = experiments_data[i]\n",
        "\n",
        "  method_x.append(method)\n",
        "  method_y_mean.append(\n",
        "      experiment_data['mean_acc_for_mask'][feature_counts[-1]]\n",
        "      / experiment_data['mean_time']\n",
        "      / divide_by\n",
        "  )\n",
        "  method_y_std.append(\n",
        "      experiment_data['std_acc_for_mask'][feature_counts[-1]]\n",
        "      / experiment_data['mean_time']\n",
        "      / divide_by\n",
        "  )\n",
        "\n",
        "method_x_pos = np.arange(len(method_x))\n",
        "\n",
        "fig, ax = plt.subplots()\n",
        "barlist = ax.bar(\n",
        "    method_x,\n",
        "    method_y_mean,\n",
        "    yerr=method_y_std,\n",
        "    align='center',\n",
        "    alpha=0.5,\n",
        "    ecolor='black',\n",
        "    capsize=10,\n",
        ")\n",
        "\n",
        "ax.set_xticks(method_x_pos)\n",
        "ax.set_xticklabels(method_x)\n",
        "ax.yaxis.grid(True)\n",
        "\n",
        "plt.ylabel('Ratio relative to Sequential attention.')\n",
        "plt.xlabel('Method')\n",
        "plt.rcParams['figure.figsize'] = (22, 9)\n",
        "plt.title(\n",
        "    'Final accuracy / Time to find the mask ratio (50 features, Activity)'\n",
        ")\n",
        "plt.show()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "last_runtime": {
        "build_target": "//learning/grp/tools/ml_python:ml_notebook",
        "kind": "private"
      },
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "1Nfxe12qK5GdSJAEiQGi2NibQGuGWmqqi",
          "timestamp": 1691164183149
        }
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
